feat: add post_create and post_generate hooks

Author: iloveitaly

polyfactory/factories/base.py

@@ -1067,6 +1067,24 @@ def process_kwargs(cls, **kwargs: Any) -> dict[str, Any]:
         for field_name, post_generator in generate_post.items():
             result[field_name] = post_generator.to_value(field_name, result)
 
+        return cls.post_generate(result)
+
+    @classmethod
+    def post_create(cls, model: T) -> None:
+        """Post-create hook. Helpful for building additional database associations or running logic which requires the
+        fully-created model.
+
+        :param model: The created model instance.
+        """
+        pass
+
+    @classmethod
+    def post_generate(cls, result: dict[str, Any]) -> dict[str, Any]:
+        """Post-generate hook. Helpful for mutating or adding additional fields right before model creation.
+
+        :param result: The kwargs that will be passed to the model.
+        :returns: The (optionally) mutated kwargs.
+        """
         return result
 
     @classmethod
@@ -1127,7 +1145,11 @@ def build(cls, *_: Any, **kwargs: Any) -> T:
         :returns: An instance of type T.
 
         """
-        return cast("T", cls.__model__(**cls.process_kwargs(**kwargs)))
+        created_model = cast("T", cls.__model__(**cls.process_kwargs(**kwargs)))
+
+        cls.post_create(created_model)
+
+        return created_model
 
     @classmethod
     def batch(cls, size: int, **kwargs: Any) -> list[T]:

polyfactory/factories/pydantic_factory.py

@@ -147,7 +147,7 @@ def from_field_info(
         min_collection_length: int | None = None,
         max_collection_length: int | None = None,
     ) -> PydanticFieldMeta:
-        """Create an instance from a pydantic field info.
+        """Create an instance from a pydantic field info. Used by `get_model_fields` to generate field list for a model.
 
         :param field_name: The name of the field.
         :param field_info: A pydantic FieldInfo instance.
@@ -473,7 +473,11 @@ def build(
 
         processed_kwargs = cls.process_kwargs(**kwargs)
 
-        return cls._create_model(kwargs["_build_context"], **processed_kwargs)
+        created_model = cls._create_model(kwargs["_build_context"], **processed_kwargs)
+
+        cls.post_create(created_model)
+
+        return created_model
 
     @classmethod
     def _get_build_context(cls, build_context: BaseBuildContext | PydanticBuildContext | None) -> PydanticBuildContext: