Deprecate the extract_prefix option to PostGeneration.

Introduces a new, call-less syntax for the @post_generation decorator.

Signed-off-by: Raphaël Barrois <raphael.barrois@polytechnique.org>

3 files changed, 136 insertions, 6 deletions

diff --git a/factory/declarations.py b/factory/declarations.py
index d5b7950..83f4d32 100644
--- a/factory/declarations.py
+++ b/factory/declarations.py
@@ -363,6 +363,11 @@ class PostGenerationDeclaration(object):
     """
 
     def __init__(self, extract_prefix=None):
+        if extract_prefix:
+            warnings.warn(
+                "The extract_prefix argument to PostGeneration declarations "
+                "is deprecated and will be removed in the future.",
+                PendingDeprecationWarning, 3)
         self.extract_prefix = extract_prefix
 
     def extract(self, name, attrs):
@@ -410,10 +415,22 @@ class PostGeneration(PostGenerationDeclaration):
         self.function(obj, create, extracted, **kwargs)
 
 
-def post_generation(extract_prefix=None):
-    def decorator(fun):
-        return PostGeneration(fun, extract_prefix=extract_prefix)
-    return decorator
+def post_generation(*args, **kwargs):
+    assert len(args) + len(kwargs) <= 1, "post_generation takes at most one argument."
+    if args and callable(args[0]):
+        # Called as @post_generation applied to a function
+        return PostGeneration(args[0])
+    else:
+        warnings.warn(
+            "The @post_generation should now be applied directly to the "
+            "function, without parameters. The @post_generation() and "
+            "@post_generation(extract_prefix='xxx') syntaxes are deprecated "
+            "and will be removed in the future; use @post_generation instead.",
+            PendingDeprecationWarning, 2)
+        extract_prefix = kwargs.get('extract_prefix')
+        def decorator(fun):
+            return PostGeneration(fun, extract_prefix=extract_prefix)
+        return decorator
 
 
 class RelatedFactory(PostGenerationDeclaration):
diff --git a/tests/test_declarations.py b/tests/test_declarations.py
index ecec244..ef7cc0f 100644
--- a/tests/test_declarations.py
+++ b/tests/test_declarations.py
@@ -122,6 +122,7 @@ class PostGenerationDeclarationTestCase(unittest.TestCase):
         self.assertEqual(extracted, 13)
         self.assertEqual(kwargs, {'bar': 42})
 
+    @tools.disable_warnings
     def test_extract_with_prefix(self):
         decl = declarations.PostGenerationDeclaration(extract_prefix='blah')
 
@@ -130,6 +131,100 @@ class PostGenerationDeclarationTestCase(unittest.TestCase):
         self.assertEqual(extracted, 42)
         self.assertEqual(kwargs, {'baz': 1})
 
+    def test_extract_prefix_deprecated(self):
+        with warnings.catch_warnings(record=True) as w:
+            __warningregistry__.clear()
+
+            warnings.simplefilter('always')
+            declarations.PostGenerationDeclaration(extract_prefix='blah')
+
+            self.assertEqual(1, len(w))
+            self.assertIn('extract_prefix', str(w[0].message))
+            self.assertIn('deprecated', str(w[0].message))
+
+    def test_decorator_simple(self):
+        call_params = []
+        @declarations.post_generation
+        def foo(*args, **kwargs):
+            call_params.append(args)
+            call_params.append(kwargs)
+
+        extracted, kwargs = foo.extract('foo',
+            {'foo': 13, 'foo__bar': 42, 'blah': 42, 'blah__baz': 1})
+        self.assertEqual(13, extracted)
+        self.assertEqual({'bar': 42}, kwargs)
+
+        # No value returned.
+        foo.call(None, False, extracted, **kwargs)
+        self.assertEqual(2, len(call_params))
+        self.assertEqual((None, False, 13), call_params[0])
+        self.assertEqual({'bar': 42}, call_params[1])
+
+    @tools.disable_warnings
+    def test_decorator_call_no_prefix(self):
+        call_params = []
+        @declarations.post_generation()
+        def foo(*args, **kwargs):
+            call_params.append(args)
+            call_params.append(kwargs)
+
+        extracted, kwargs = foo.extract('foo',
+            {'foo': 13, 'foo__bar': 42, 'blah': 42, 'blah__baz': 1})
+        self.assertEqual(13, extracted)
+        self.assertEqual({'bar': 42}, kwargs)
+
+        # No value returned.
+        foo.call(None, False, extracted, **kwargs)
+        self.assertEqual(2, len(call_params))
+        self.assertEqual((None, False, 13), call_params[0])
+        self.assertEqual({'bar': 42}, call_params[1])
+
+    @tools.disable_warnings
+    def test_decorator_extract_prefix(self):
+        call_params = []
+        @declarations.post_generation(extract_prefix='blah')
+        def foo(*args, **kwargs):
+            call_params.append(args)
+            call_params.append(kwargs)
+
+        extracted, kwargs = foo.extract('foo',
+            {'foo': 13, 'foo__bar': 42, 'blah': 42, 'blah__baz': 1})
+        self.assertEqual(42, extracted)
+        self.assertEqual({'baz': 1}, kwargs)
+
+        # No value returned.
+        foo.call(None, False, extracted, **kwargs)
+        self.assertEqual(2, len(call_params))
+        self.assertEqual((None, False, 42), call_params[0])
+        self.assertEqual({'baz': 1}, call_params[1])
+
+    def test_decorator_call_no_prefix_deprecated(self):
+        with warnings.catch_warnings(record=True) as w:
+            __warningregistry__.clear()
+
+            warnings.simplefilter('always')
+            @declarations.post_generation()
+            def foo(*args, **kwargs):
+                pass
+
+            self.assertEqual(1, len(w))
+            self.assertIn('post_generation', str(w[0].message))
+            self.assertIn('deprecated', str(w[0].message))
+
+    def test_decorator_call_with_prefix_deprecated(self):
+        with warnings.catch_warnings(record=True) as w:
+            __warningregistry__.clear()
+
+            warnings.simplefilter('always')
+            @declarations.post_generation(extract_prefix='blah')
+            def foo(*args, **kwargs):
+                pass
+
+            # 2 warnings: decorator with brackets, and extract_prefix.
+            self.assertEqual(2, len(w))
+            self.assertIn('post_generation', str(w[0].message))
+            self.assertIn('deprecated', str(w[0].message))
+
 
 class SubFactoryTestCase(unittest.TestCase):
     def test_lazyness(self):
@@ -174,7 +269,6 @@ class CircularSubFactoryTestCase(unittest.TestCase):
             self.assertIn('CircularSubFactory', str(w[0].message))
             self.assertIn('deprecated', str(w[0].message))
 
-
     @tools.disable_warnings
     def test_lazyness(self):
         f = declarations.CircularSubFactory('factory.declarations', 'Sequence', x=3)
diff --git a/tests/test_using.py b/tests/test_using.py
index 2bab8c0..fd8befb 100644
--- a/tests/test_using.py
+++ b/tests/test_using.py
@@ -1173,6 +1173,25 @@ class PostGenerationTestCase(unittest.TestCase):
 
             one = 1
 
+            @factory.post_generation
+            def incr_one(self, _create, _increment):
+                self.one += 1
+
+        obj = TestObjectFactory.build()
+        self.assertEqual(2, obj.one)
+        self.assertFalse(hasattr(obj, 'incr_one'))
+
+        obj = TestObjectFactory.build(one=2)
+        self.assertEqual(3, obj.one)
+        self.assertFalse(hasattr(obj, 'incr_one'))
+
+    @tools.disable_warnings
+    def test_post_generation_calling(self):
+        class TestObjectFactory(factory.Factory):
+            FACTORY_FOR = TestObject
+
+            one = 1
+
             @factory.post_generation()
             def incr_one(self, _create, _increment):
                 self.one += 1
@@ -1191,7 +1210,7 @@ class PostGenerationTestCase(unittest.TestCase):
 
             one = 1
 
-            @factory.post_generation()
+            @factory.post_generation
             def incr_one(self, _create, increment=1):
                 self.one += increment