fix: Dataclass field type not used correctly (#371)

Co-authored-by: Andrew Truong <[email protected]>

Author: anthonyjgraff

polyfactory/factories/dataclass_factory.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 from dataclasses import MISSING, fields, is_dataclass
-from typing import Any, Generic
-
-from typing_extensions import TypeGuard, get_type_hints
+from typing import TYPE_CHECKING, Any, ForwardRef, Generic
 
 from polyfactory.factories.base import BaseFactory, T
 from polyfactory.field_meta import FieldMeta, Null
+from polyfactory.utils.helpers import evaluate_forwardref
+
+if TYPE_CHECKING:
+    from typing_extensions import TypeGuard
 
 
 class DataclassFactory(Generic[T], BaseFactory[T]):
@@ -33,8 +35,6 @@ def get_model_fields(cls) -> list["FieldMeta"]:
         """
         fields_meta: list["FieldMeta"] = []
 
-        model_type_hints = get_type_hints(cls.__model__, include_extras=True)
-
         for field in fields(cls.__model__):  # type: ignore[arg-type]
             if field.default_factory and field.default_factory is not MISSING:
                 default_value = field.default_factory()
@@ -43,9 +43,16 @@ def get_model_fields(cls) -> list["FieldMeta"]:
             else:
                 default_value = Null
 
+            if isinstance(field.type, ForwardRef):
+                annotation = evaluate_forwardref(field.type)  # type: ignore[unreachable]
+            elif isinstance(field.type, str):
+                annotation = evaluate_forwardref(ForwardRef(field.type))  # type: ignore[unreachable]
+            else:
+                annotation = field.type
+
             fields_meta.append(
                 FieldMeta.from_type(
-                    annotation=model_type_hints[field.name],
+                    annotation=annotation,
                     name=field.name,
                     default=default_value,
                     random=cls.__random__,

polyfactory/utils/helpers.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import sys
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ForwardRef
 
 from typing_extensions import get_args, get_origin
 
@@ -130,3 +130,17 @@ def normalize_annotation(annotation: Any, random: Random) -> Any:
         return origin[args] if origin is not type else annotation
 
     return origin
+
+
+def evaluate_forwardref(ref: ForwardRef) -> Any:
+    """Evaluate ForwardRef to get type annotation
+
+    :param ref: A ForwardRef object
+
+    :returns: A type annotation
+
+    """
+    if sys.version_info < (3, 9):
+        return ref._evaluate(globals(), locals())
+
+    return ref._evaluate(globals(), locals(), frozenset())

tests/test_dataclass_factory.py

@@ -55,6 +55,27 @@ class MyFactory(DataclassFactory[PydanticDC]):
     assert result.constrained_field >= 100
 
 
+def test_factory_sqlalchemy_dc() -> None:
+    from sqlalchemy.orm import DeclarativeBase, Mapped, MappedAsDataclass, mapped_column
+
+    class SqlAlchemyDCBase(MappedAsDataclass, DeclarativeBase):
+        pass
+
+    class SqlAlchemyDC(SqlAlchemyDCBase):
+        __tablename__ = "foo"
+
+        id: Mapped[int] = mapped_column(primary_key=True)
+        name: Mapped[str]
+
+    class SqlAlchemyDCFactory(DataclassFactory[SqlAlchemyDC]):
+        __model__ = SqlAlchemyDC
+
+    result = SqlAlchemyDCFactory.build()
+
+    assert isinstance(result.id, int)
+    assert isinstance(result.name, str)
+
+
 def test_vanilla_dc_with_embedded_model() -> None:
     @vanilla_dataclass
     class VanillaDC:
@@ -168,6 +189,36 @@ class MyFactory(DataclassFactory[example]):  # type:ignore[valid-type]
     assert MyFactory.process_kwargs() == {"foo": ANY}
 
 
+def test_sqlalchemy_dc_factory_with_future_annotations(create_module: Callable[[str], ModuleType]) -> None:
+    module = create_module(
+        """
+from __future__ import annotations
+
+from sqlalchemy.orm import DeclarativeBase, Mapped, MappedAsDataclass, mapped_column
+
+class SqlAlchemyDCBase(MappedAsDataclass, DeclarativeBase):  # type: ignore
+    pass
+
+class SqlAlchemyDC(SqlAlchemyDCBase):
+    __tablename__ = "foo"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    name: Mapped[str]
+
+"""
+    )
+
+    SqlAlchemyDC: type = module.SqlAlchemyDC
+    assert SqlAlchemyDC.__annotations__ == {"id": "Mapped[int]", "name": "Mapped[str]"}
+
+    class SqlAlchemyDCFactory(DataclassFactory[SqlAlchemyDC]):  # type:ignore[valid-type]
+        __model__ = SqlAlchemyDC
+
+    result = SqlAlchemyDCFactory.build()
+    assert isinstance(result.id, int)
+    assert isinstance(result.name, str)
+
+
 def test_variable_length_tuple_generation__many_type_args() -> None:
     @vanilla_dataclass
     class VanillaDC: