Add a 'after post_generation' hook to Factory.

Use it in DjangoModelFactory to save objects again if a post_generation hook ran.

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

2 files changed, 24 insertions, 2 deletions

diff --git a/factory/base.py b/factory/base.py
index 28d7cdb..3ebc746 100644
--- a/factory/base.py
+++ b/factory/base.py
@@ -616,12 +616,27 @@ class Factory(BaseFactory):
         obj = cls._prepare(create, **attrs)
 
         # Handle post-generation attributes
+        results = {}
         for name, decl in sorted(postgen_declarations.items()):
             extracted, extracted_kwargs = postgen_attributes[name]
-            decl.call(obj, create, extracted, **extracted_kwargs)
+            results[name] = decl.call(obj, create, extracted, **extracted_kwargs)
+
+        cls._after_postgeneration(obj, create, results)
+
         return obj
 
     @classmethod
+    def _after_postgeneration(cls, obj, create, results=None):
+        """Hook called after post-generation declarations have been handled.
+
+        Args:
+            obj (object): the generated object
+            create (bool): whether the strategy was 'build' or 'create'
+            results (dict or None): result of post-generation declarations
+        """
+        pass
+
+    @classmethod
     def build(cls, **kwargs):
         attrs = cls.attributes(create=False, extra=kwargs)
         return cls._generate(False, attrs)
@@ -657,6 +672,13 @@ class DjangoModelFactory(Factory):
         """Create an instance of the model, and save it to the database."""
         return target_class._default_manager.create(*args, **kwargs)
 
+    @classmethod
+    def _after_postgeneration(cls, obj, create, results=None):
+        """Save again the instance if creating and at least one hook ran."""
+        if create and results:
+            # Some post-generation hooks ran, and may have modified us.
+            obj.save()
+
 
 class MogoFactory(Factory):
     """Factory for mogo objects."""
diff --git a/factory/declarations.py b/factory/declarations.py
index d3d7659..366c2c8 100644
--- a/factory/declarations.py
+++ b/factory/declarations.py
@@ -412,7 +412,7 @@ class PostGeneration(PostGenerationDeclaration):
         self.function = function
 
     def call(self, obj, create, extracted=None, **kwargs):
-        self.function(obj, create, extracted, **kwargs)
+        return self.function(obj, create, extracted, **kwargs)
 
 
 def post_generation(*args, **kwargs):