diff --git a/sqlalchemy_utils/decorators.py b/sqlalchemy_utils/decorators.py index e10c2fc..2719ca5 100644 --- a/sqlalchemy_utils/decorators.py +++ b/sqlalchemy_utils/decorators.py @@ -21,6 +21,13 @@ class AttributeValueGenerator(object): self.reset() def reset(self): + if ( + hasattr(self, 'listeners_registered') and + self.listeners_registered + ): + for args in self.listener_args: + sa.event.remove(*args) + self.listeners_registered = False self.generator_registry = defaultdict(list) @@ -28,12 +35,8 @@ class AttributeValueGenerator(object): def wrapper(self, *args, **kwargs): return func(self, *args, **kwargs) - if isinstance(attr, six.string_types) and '.' in attr: - parts = attr.split('.') - self.generator_registry[parts[0]].append(wrapper) - wrapper.__generates__ = parts[1] - elif isinstance(attr, sa.orm.attributes.InstrumentedAttribute): - self.generator_registry[attr.class_.__name__].append(wrapper) + if isinstance(attr, sa.orm.attributes.InstrumentedAttribute): + self.generator_registry[attr.class_].append(wrapper) wrapper.__generates__ = attr else: wrapper.__generates__ = attr @@ -46,18 +49,22 @@ class AttributeValueGenerator(object): self.listeners_registered = True def update_generator_registry(self, mapper, class_): + """ + Adds generator functions to generator_registry. + """ for value in class_.__dict__.values(): if hasattr(value, '__generates__'): - self.generator_registry[class_.__name__].append(value) + self.generator_registry[class_].append(value) def update_generated_properties(self, session, ctx, instances): for obj in itertools.chain(session.new, session.dirty): - class_ = obj.__class__.__name__ + class_ = obj.__class__ if class_ in self.generator_registry: for func in self.generator_registry[class_]: attr = func.__generates__ if not isinstance(attr, six.string_types): attr = attr.name + print obj, attr setattr(obj, attr, func(obj)) @@ -120,16 +127,6 @@ def generates(attr, generator=generator): @generates(Article.slug) def _create_article_slug(article): return article.name.lower().replace(' ', '-') - - - Or with lazy evaluated string argument: - - :: - - - @generates('Article.slug') - def _create_article_slug(article): - return article.name.lower().replace(' ', '-') """ generator.register_listeners() diff --git a/tests/test_generates.py b/tests/test_generates.py index a8a7c38..81a8cd7 100644 --- a/tests/test_generates.py +++ b/tests/test_generates.py @@ -43,21 +43,6 @@ class TestGeneratesWithBoundMethodAndStringArg(GeneratesTestCase): self.Article = Article -class TestGeneratesWithFunctionAndStringArg(GeneratesTestCase): - def create_models(self): - class Article(self.Base): - __tablename__ = 'article' - id = sa.Column(sa.Integer, primary_key=True) - name = sa.Column(sa.Unicode(255)) - slug = sa.Column(sa.Unicode(255)) - - @generates('Article.slug') - def _create_article_slug(self): - return self.name.lower().replace(' ', '-') - - self.Article = Article - - class TestGeneratesWithFunctionAndClassVariableArg(GeneratesTestCase): def create_models(self): class Article(self.Base):