Skip to content

Commit e4a956c

Browse files
committed
feat: add post_create and post_generate hooks
1 parent 02fed72 commit e4a956c

2 files changed

Lines changed: 29 additions & 3 deletions

File tree

polyfactory/factories/base.py

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1071,6 +1071,24 @@ def process_kwargs(cls, **kwargs: Any) -> dict[str, Any]:
10711071
for field_name, post_generator in generate_post.items():
10721072
result[field_name] = post_generator.to_value(field_name, result)
10731073

1074+
return cls.post_generate(result)
1075+
1076+
@classmethod
1077+
def post_create(cls, model: T) -> None:
1078+
"""Post-create hook. Helpful for building additional database associations or running logic which requires the
1079+
fully-created model.
1080+
1081+
:param model: The created model instance.
1082+
"""
1083+
pass
1084+
1085+
@classmethod
1086+
def post_generate(cls, result: dict[str, Any]) -> dict[str, Any]:
1087+
"""Post-generate hook. Helpful for mutating or adding additional fields right before model creation.
1088+
1089+
:param result: The kwargs that will be passed to the model.
1090+
:returns: The (optionally) mutated kwargs.
1091+
"""
10741092
return result
10751093

10761094
@classmethod
@@ -1131,7 +1149,11 @@ def build(cls, *_: Any, **kwargs: Any) -> T:
11311149
:returns: An instance of type T.
11321150
11331151
"""
1134-
return cast("T", cls.__model__(**cls.process_kwargs(**kwargs)))
1152+
created_model = cast("T", cls.__model__(**cls.process_kwargs(**kwargs)))
1153+
1154+
cls.post_create(created_model)
1155+
1156+
return created_model
11351157

11361158
@classmethod
11371159
def batch(cls, size: int, **kwargs: Any) -> list[T]:

polyfactory/factories/pydantic_factory.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,7 @@ def from_field_info(
147147
min_collection_length: int | None = None,
148148
max_collection_length: int | None = None,
149149
) -> PydanticFieldMeta:
150-
"""Create an instance from a pydantic field info.
150+
"""Create an instance from a pydantic field info. Used by `get_model_fields` to generate field list for a model.
151151
152152
:param field_name: The name of the field.
153153
:param field_info: A pydantic FieldInfo instance.
@@ -517,7 +517,11 @@ def build(
517517

518518
processed_kwargs = cls.process_kwargs(**kwargs)
519519

520-
return cls._create_model(kwargs["_build_context"], **processed_kwargs)
520+
created_model = cls._create_model(kwargs["_build_context"], **processed_kwargs)
521+
522+
cls.post_create(created_model)
523+
524+
return created_model
521525

522526
@classmethod
523527
def _get_build_context(cls, build_context: BaseBuildContext | PydanticBuildContext | None) -> PydanticBuildContext:

0 commit comments

Comments
 (0)