From 002b5b69cd0d437be298414949ef59bf14d57df6 Mon Sep 17 00:00:00 2001 From: Konsta Vesterinen Date: Wed, 23 Oct 2013 15:33:15 +0300 Subject: [PATCH] Fixed generates decorator --- sqlalchemy_utils/decorators.py | 3 +++ tests/test_generates.py | 4 ++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/sqlalchemy_utils/decorators.py b/sqlalchemy_utils/decorators.py index 378053b..decb5eb 100644 --- a/sqlalchemy_utils/decorators.py +++ b/sqlalchemy_utils/decorators.py @@ -18,6 +18,9 @@ def generates(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 diff --git a/tests/test_generates.py b/tests/test_generates.py index e41bf5a..9d02427 100644 --- a/tests/test_generates.py +++ b/tests/test_generates.py @@ -5,8 +5,8 @@ from tests import TestCase class GeneratesTestCase(TestCase): - def setup_method(self, method): - TestCase.setup_method(self, method) + def teardown_method(self, method): + TestCase.teardown_method(self, method) decorators.generator_registry = defaultdict(list) def test_generates_value_before_flush(self):