62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
from collections import defaultdict
|
|
import itertools
|
|
import sqlalchemy as sa
|
|
import six
|
|
|
|
generator_registry = defaultdict(list)
|
|
listeners_registered = False
|
|
|
|
|
|
def generates(attr):
|
|
register_listeners()
|
|
|
|
def wraps(func):
|
|
def wrapper(self, *args, **kwargs):
|
|
return func(self, *args, **kwargs)
|
|
|
|
if isinstance(attr, six.string_types) and '.' in attr:
|
|
parts = attr.split('.')
|
|
generator_registry[parts[0]].append(wrapper)
|
|
wrapper.__generates__ = parts[1]
|
|
elif isinstance(attr, sa.orm.attributes.InstrumentedAttribute):
|
|
generator_registry[attr.class_.__name__].append(wrapper)
|
|
wrapper.__generates__ = attr
|
|
else:
|
|
wrapper.__generates__ = attr
|
|
return wrapper
|
|
return wraps
|
|
|
|
|
|
def register_listeners():
|
|
global listeners_registered
|
|
|
|
if not listeners_registered:
|
|
sa.event.listen(
|
|
sa.orm.mapper,
|
|
'mapper_configured',
|
|
update_generator_registry
|
|
)
|
|
sa.event.listen(
|
|
sa.orm.session.Session,
|
|
'before_flush',
|
|
update_generated_properties
|
|
)
|
|
listeners_registered = True
|
|
|
|
|
|
def update_generator_registry(mapper, class_):
|
|
for value in class_.__dict__.values():
|
|
if hasattr(value, '__generates__'):
|
|
generator_registry[class_.__name__].append(value)
|
|
|
|
|
|
def update_generated_properties(session, ctx, instances):
|
|
for obj in itertools.chain(session.new, session.dirty):
|
|
class_ = obj.__class__.__name__
|
|
if class_ in generator_registry:
|
|
for func in generator_registry[class_]:
|
|
attr = func.__generates__
|
|
if not isinstance(attr, six.string_types):
|
|
attr = attr.name
|
|
setattr(obj, attr, func(obj))
|