Repository: litestar-org/polyfactory
Pull request #680 (2025-03-29 07:10:19)

feat: support new type and annotated mapping
<!--
By submitting this pull request, you agree to:
- follow [Litestar's Code of Conduct](https://github.com/litestar-org/.github/blob/main/CODE_OF_CONDUCT.md)
- follow [Litestar's contribution guidelines](https://github.com/litestar-org/.github/blob/main/CONTRIBUTING.md)
- follow the [PSFs's Code of Conduct](https://www.python.org/psf/conduct/)
-->
## Description

- Allow applying provider map to apply to annotated or new types
- This requires updating `FieldMeta` to preserve original type. I do not think this is backwards breaking but happy to hold off to v3 to support this 

<!--
Please add in issue numbers this pull request will close, if applicable
Examples: Fixes #4321 or Closes #1234

Ensure you are using a supported keyword to properly link an issue:
https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword
-->
## Closes

- Closes https://github.com/litestar-org/polyfactory/issues/348

github-actions[bot]: Documentation preview will be available shortly at https://litestar-org.github.io/polyfactory-docs-preview/680

Patch:
diff --git a/polyfactory/factories/base.py b/polyfactory/factories/base.py
index aef5cb8..0905803 100644
--- a/polyfactory/factories/base.py
+++ b/polyfactory/factories/base.py
@@ -835,7 +835,8 @@ class BaseFactory(ABC, Generic[T]):
                 build_context=build_context,
             )
 
-        if provider := cls.get_provider_map().get(unwrapped_annotation):
+        provider_map = cls.get_provider_map()
+        if provider := (provider_map.get(field_meta.annotation) or provider_map.get(unwrapped_annotation)):
             return provider()
 
         if isinstance(unwrapped_annotation, TypeVar):
@@ -914,7 +915,10 @@ class BaseFactory(ABC, Generic[T]):
 
                 yield handle_collection_type_coverage(child_meta, origin, cls, build_context=build_context)
 
-            elif provider := cls.get_provider_map().get(unwrapped_annotation):
+            elif provider := (
+                (provider_map := cls.get_provider_map()).get(field_meta.annotation)
+                or provider_map.get(unwrapped_annotation)
+            ):
                 yield CoverageContainerCallable(provider)
 
             elif isinstance(unwrapped_annotation, TypeVar):
diff --git a/polyfactory/field_meta.py b/polyfactory/field_meta.py
index 553631b..c27610f 100644
--- a/polyfactory/field_meta.py
+++ b/polyfactory/field_meta.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import functools
 from dataclasses import asdict
 from typing import TYPE_CHECKING, Any, Hashable, Literal, Mapping, Pattern, TypedDict, cast
 
@@ -60,7 +61,7 @@ class Constraints(TypedDict):
 class FieldMeta:
     """Factory field metadata container. This class is used to store the data about a field of a factory's model."""
 
-    __slots__ = ("annotation", "children", "constraints", "default", "name", "random")
+    __slots__ = ("__dict__", "annotation", "children", "constraints", "default", "name", "random")
 
     annotation: Any
     random: Random
@@ -88,15 +89,18 @@ class FieldMeta:
         self.name = name
         self.constraints = constraints
 
-    @property
+    @functools.cached_property
     def type_args(self) -> tuple[Any, ...]:
         """Return the normalized type args of the annotation, if any.
 
         :returns: a tuple of types.
         """
+        annotation = self.annotation
+        if is_annotated(annotation):
+            annotation = get_args(self.annotation)[0]
         return tuple(
             TYPE_MAPPING.get(arg, arg) if isinstance(arg, Hashable) else arg  # pyright: ignore[reportCallIssue,reportArgumentType]
-            for arg in get_args(self.annotation)
+            for arg in get_args(annotation)
         )
 
     @classmethod
@@ -140,9 +144,7 @@ class FieldMeta:
             metadata = cls.get_constraints_metadata(annotation)
             constraints = cls.parse_constraints(metadata)
 
-        if annotated:
-            annotation = get_args(annotation)[0]
-        elif (origin := get_origin(annotation)) and origin in TYPE_MAPPING:  # pragma: no cover
+        if not annotated and (origin := get_origin(annotation)) and origin in TYPE_MAPPING:
             container = TYPE_MAPPING[origin]
             annotation = container[get_args(annotation)]  # type: ignore[index]
 


Test patch:
diff --git a/tests/sqlalchemy_factory/test_sqlalchemy_factory_common.py b/tests/sqlalchemy_factory/test_sqlalchemy_factory_common.py
index f8b3bf8..37feccf 100644
--- a/tests/sqlalchemy_factory/test_sqlalchemy_factory_common.py
+++ b/tests/sqlalchemy_factory/test_sqlalchemy_factory_common.py
@@ -1,7 +1,7 @@
 from datetime import datetime
 from decimal import Decimal
 from enum import Enum
-from typing import Any, Callable
+from typing import Any, Callable, get_args
 from uuid import UUID
 
 import pytest
@@ -432,7 +432,7 @@ def test_numeric_field(numeric: Numeric) -> None:
     class NumericModelFactory(SQLAlchemyFactory[NumericModel]): ...
 
     result = NumericModelFactory.get_model_fields()[1]
-    assert result.annotation is Decimal
+    assert result.annotation is Decimal or get_args(result.annotation)[0] is Decimal
     if constraints := result.constraints:
         assert constraints.get("max_digits") == numeric.precision
         assert constraints.get("decimal_places") == numeric.scale
diff --git a/tests/test_base_factories.py b/tests/test_base_factories.py
index 4da6f61..de8c96b 100644
--- a/tests/test_base_factories.py
+++ b/tests/test_base_factories.py
@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, NewType
 
 import pytest
+from typing_extensions import Annotated
 
 from pydantic.main import BaseModel
 
@@ -129,3 +130,32 @@ def test_abstract_classes_are_ignored() -> None:
 
     with pytest.raises(ParameterException, match="Unsupported type: "):
         DataclassFactory.create_factory(Model).build()
+
+
+AnnotatedType = Annotated[str, "foo"]
+CustomNewType = NewType("CustomNewType", str)
+
+
+def test_annotated_types_used_as_provider() -> None:
+    @dataclass
+    class Model:
+        a: AnnotatedType
+        b: CustomNewType
+
+    class Factory(DataclassFactory[Model]):
+        @classmethod
+        def get_provider_map(cls) -> Dict[Any, Callable[[], Any]]:
+            return {
+                **super().get_provider_map(),
+                AnnotatedType: lambda: "foo",
+                CustomNewType: lambda: "bar",
+            }
+
+    result = Factory.build()
+    assert result.a == "foo"
+    assert result.b == "bar"
+
+    result_coverage = list(Factory.coverage())
+    assert len(result_coverage) == 1
+    assert result_coverage[0].a == "foo"
+    assert result_coverage[0].b == "bar"