feat: improve create factory typing (#657)

Author: adhtruong

polyfactory/factories/base.py

@@ -39,11 +39,12 @@
     TypedDict,
     TypeVar,
     cast,
+    overload,
 )
 from uuid import UUID
 
 from faker import Faker
-from typing_extensions import get_args, get_origin, get_original_bases
+from typing_extensions import Self, get_args, get_origin, get_original_bases
 
 from polyfactory.constants import (
     DEFAULT_RANDOM,
@@ -89,6 +90,7 @@
 
 
 T = TypeVar("T")
+U = TypeVar("U")
 F = TypeVar("F", bound="BaseFactory[Any]")
 
 
@@ -372,7 +374,7 @@ def _get_config(cls) -> dict[str, Any]:
         }
 
     @classmethod
-    def _get_or_create_factory(cls, model: type) -> type[BaseFactory[Any]]:
+    def _get_or_create_factory(cls, model: type[U]) -> type[BaseFactory[U]]:
         """Get a factory from registered factories or generate a factory dynamically.
 
         :param model: A model type.
@@ -553,13 +555,31 @@ def _create_generic_fn() -> Callable:
             **(cls._extra_providers or {}),
         }
 
+    @overload
     @classmethod
     def create_factory(
-        cls: type[F],
-        model: type[T] | None = None,
+        cls,
+        model: None = None,
+        bases: tuple[type[BaseFactory[Any]], ...] | None = None,
+        **kwargs: Any,
+    ) -> type[Self]: ...
+
+    @overload
+    @classmethod
+    def create_factory(
+        cls,
+        model: type[U],
         bases: tuple[type[BaseFactory[Any]], ...] | None = None,
         **kwargs: Any,
-    ) -> type[F]:
+    ) -> type[BaseFactory[U]]: ...
+
+    @classmethod
+    def create_factory(
+        cls,
+        model: type[U] | None = None,
+        bases: tuple[type[BaseFactory[Any]], ...] | None = None,
+        **kwargs: Any,
+    ) -> type[Self | BaseFactory[U]]:
         """Generate a factory for the given type dynamically.
 
         :param model: A type to model. Defaults to current factory __model__ if any.
@@ -572,12 +592,13 @@ def create_factory(
         """
         if model is None:
             try:
-                model = cls.__model__
+                model = cls.__model__  # pyright: ignore[reportAssignmentType]
             except AttributeError as ex:
                 msg = "A 'model' argument is required when creating a new factory from a base one"
                 raise TypeError(msg) from ex
+
         return cast(
-            "Type[F]",
+            "Type[Self]",
             type(
                 f"{model.__name__}Factory",  # pyright: ignore[reportOptionalMemberAccess]
                 (*(bases or ()), cls),
@@ -1081,7 +1102,7 @@ def process_kwargs_coverage(cls, **kwargs: Any) -> abc.Iterable[dict[str, Any]]:
             yield resolved
 
     @classmethod
-    def build(cls, **kwargs: Any) -> T:
+    def build(cls, *_: Any, **kwargs: Any) -> T:
         """Build an instance of the factory's __model__
 
         :param kwargs: Any kwargs. If field names are set in kwargs, their values will be used.

polyfactory/factories/pydantic_factory.py

@@ -97,7 +97,9 @@
 
     from typing_extensions import NotRequired, TypeGuard
 
-T = TypeVar("T", bound="BaseModelV1 | BaseModelV2")  # pyright: ignore[reportInvalidTypeForm]
+    from pydantic import BaseModel
+
+T = TypeVar("T", bound="BaseModel")
 
 _IS_PYDANTIC_V1 = VERSION.startswith("1")
 
@@ -465,7 +467,8 @@ def build(
 
         if "_build_context" not in kwargs:
             kwargs["_build_context"] = PydanticBuildContext(
-                seen_models=set(), factory_use_construct=factory_use_construct
+                seen_models=set(),
+                factory_use_construct=factory_use_construct,
             )
 
         processed_kwargs = cls.process_kwargs(**kwargs)
@@ -502,8 +505,8 @@ def _create_model(cls, _build_context: PydanticBuildContext, **kwargs: Any) -> T
         if cls._get_build_context(_build_context).get("factory_use_construct"):
             if _is_pydantic_v1_model(cls.__model__):
                 return cls.__model__.construct(**kwargs)  # type: ignore[return-value]
-            return cls.__model__.model_construct(**kwargs)  # type: ignore[return-value]
-        return cls.__model__(**kwargs)  # type: ignore[return-value]
+            return cls.__model__.model_construct(**kwargs)
+        return cls.__model__(**kwargs)
 
     @classmethod
     def coverage(cls, factory_use_construct: bool = False, **kwargs: Any) -> abc.Iterator[T]:

tests/test_new_types.py

@@ -160,6 +160,6 @@ class A:
 """,
     )
 
-    factory = DataclassFactory.create_factory(module.A)  # type: ignore[var-annotated]
+    factory = DataclassFactory.create_factory(module.A)
     result = factory.build()
     assert result.field in {"a", "b"}

tests/test_pydantic_v1_v2.py

@@ -49,7 +49,7 @@ def test_build_v1_with_contrained_fields() -> None:
     ConstrainedInt = Annotated[int, Field(ge=100, le=200)]
     ConstrainedStr = Annotated[str, Field(min_length=1, max_length=3)]
 
-    class Foo(pydantic.v1.BaseModel):  # pyright: ignore[reportGeneralTypeIssues]
+    class Foo(BaseModelV1):
         a: ConstrainedInt
         b: ConstrainedStr
         c: Union[ConstrainedInt, ConstrainedStr]
@@ -58,7 +58,7 @@ class Foo(pydantic.v1.BaseModel):  # pyright: ignore[reportGeneralTypeIssues]
         f: List[ConstrainedInt]
         g: Dict[ConstrainedInt, ConstrainedStr]
 
-    ModelFactory.create_factory(Foo).build()  # type: ignore[type-var]
+    ModelFactory.create_factory(Foo).build()
 
 
 def test_build_v2_with_contrained_fields() -> None:

tests/test_recursive_models.py

@@ -53,7 +53,7 @@ class PydanticNode(BaseModel):
 
 @pytest.mark.parametrize("factory_use_construct", (True, False))
 def test_recursive_pydantic_models(factory_use_construct: bool) -> None:
-    factory = ModelFactory[PydanticNode].create_factory(PydanticNode)
+    factory = ModelFactory.create_factory(PydanticNode)
 
     result = factory.build(factory_use_construct)
     assert result.child is _Sentinel, "Default is not used"  # type: ignore[comparison-overlap]