feat(sqlalchemy-factory): Add support for SQLAlchemy custom types

docs/examples/library_factories/sqlalchemy_factory/test_example_4.py

@@ -0,0 +1,31 @@
+import datetime
+from typing import Any
+
+from sqlalchemy import DateTime, types
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+
+from polyfactory.factories.sqlalchemy_factory import SQLAlchemyFactory
+
+
+class TZAwareDateTime(types.TypeDecorator):
+    impl = DateTime(timezone=True)
+
+
+class Base(DeclarativeBase):
+    ...
+
+
+class Author(Base):
+    __tablename__ = "authors"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    first_publication_at: Mapped[Any] = mapped_column(type_=TZAwareDateTime(), nullable=False)
+
+
+class AuthorFactory(SQLAlchemyFactory[Author]):
+    __model__ = Author
+
+
+def test_sqla_type_decorator_custom_type_factory() -> None:
+    author = AuthorFactory.build()
+    assert isinstance(author.first_publication_at, datetime.datetime)

docs/examples/library_factories/sqlalchemy_factory/test_example_5.py

@@ -0,0 +1,35 @@
+import datetime
+from typing import Any, Callable, Dict
+
+from sqlalchemy import types
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+
+from polyfactory.factories.sqlalchemy_factory import SQLAlchemyFactory
+
+
+class CustomType(types.UserDefinedType):
+    ...
+
+
+class Base(DeclarativeBase):
+    ...
+
+
+class Author(Base):
+    __tablename__ = "authors"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    first_publication_at: Mapped[Any] = mapped_column(type_=CustomType())
+
+
+class AuthorFactory(SQLAlchemyFactory[Author]):
+    __model__ = Author
+
+    @classmethod
+    def get_sqlalchemy_types(cls) -> Dict[Any, Callable[[], Any]]:
+        return {**super().get_sqlalchemy_types(), CustomType: lambda: cls.__faker__.date_time()}
+
+
+def test_sqla_user_defined_type_custom_type_factory() -> None:
+    author = AuthorFactory.build()
+    assert isinstance(author.first_publication_at, datetime.datetime)

docs/usage/library_factories/sqlalchemy_factory.rst

@@ -36,6 +36,29 @@ By default, this will add generated models to the session and then commit. This
 
 Similarly for ``__async_session__`` and ``create_async``.
 
+SQLAlchemy custom types
+------------------------------
+
+TypeDecorator based
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Type will be inferred from SQLAlchemy type defined in ``impl``.
+
+.. literalinclude:: /examples/library_factories/sqlalchemy_factory/test_example_4.py
+    :caption: Custom types with TypeDecorator
+    :language: python
+
+
+UserDefinedType based
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+``get_sqlalchemy_types`` classmethod needs to be overridden to provide factory function for custom type.
+
+.. literalinclude:: /examples/library_factories/sqlalchemy_factory/test_example_5.py
+    :caption: Custom types with UserDefinedType
+    :language: python
+
+More info on `SQLAlchemy Custom Types <https://docs.sqlalchemy.org/en/20/core/custom_types.html>`_.
 
 API reference
 ------------------------------

polyfactory/factories/sqlalchemy_factory.py

@@ -3,7 +3,7 @@
 from datetime import date, datetime
 from typing import TYPE_CHECKING, Any, Callable, ClassVar, Generic, List, TypeVar, Union
 
-from polyfactory.exceptions import MissingDependencyException
+from polyfactory.exceptions import MissingDependencyException, ParameterException
 from polyfactory.factories.base import BaseFactory
 from polyfactory.field_meta import FieldMeta
 from polyfactory.persistence import AsyncPersistenceProtocol, SyncPersistenceProtocol
@@ -118,6 +118,17 @@ def get_type_from_column(cls, column: Column) -> type:
             annotation = column_type
         elif issubclass(column_type, types.ARRAY):
             annotation = List[column.type.item_type.python_type]  # type: ignore[assignment,name-defined]
+        elif issubclass(column_type, types.TypeDecorator) and isinstance(
+            python_type := column_type.impl.python_type,
+            type,
+        ):
+            annotation = python_type
+        elif issubclass(column_type, types.UserDefinedType):
+            parameter_exc_msg = (
+                f"User defined type detected (subclass of {types.UserDefinedType}). "
+                "Override get_sqlalchemy_types to provide factory function."
+            )
+            raise ParameterException(parameter_exc_msg)
         else:
             annotation = column.type.python_type
 

tests/sqlalchemy_factory/test_sqlalchemy_factory_v2.py

@@ -1,9 +1,11 @@
+import datetime
 from enum import Enum
-from typing import Any, List
+from typing import Any, Callable, Dict, List
 
 import pytest
-from sqlalchemy import ForeignKey, __version__, orm, types
+from sqlalchemy import ForeignKey, __version__, orm, sql, types
 
+from polyfactory.exceptions import ParameterException
 from polyfactory.factories.sqlalchemy_factory import SQLAlchemyFactory
 
 if __version__.startswith("1"):
@@ -85,3 +87,95 @@ class ModelFactory(SQLAlchemyFactory[Model]):
 
     instance = ModelFactory.build()
     assert instance.overridden is not None
+
+
+@pytest.mark.parametrize(
+    "impl_",
+    (
+        sql.sqltypes.BigInteger(),
+        sql.sqltypes.Boolean(),
+        sql.sqltypes.Date(),
+        sql.sqltypes.DateTime(),
+        sql.sqltypes.Double(),
+        sql.sqltypes.Enum(),
+        sql.sqltypes.Float(),
+        sql.sqltypes.Integer(),
+        sql.sqltypes.Interval(),
+        sql.sqltypes.LargeBinary(),
+        sql.sqltypes.MatchType(),
+        sql.sqltypes.Numeric(),
+        sql.sqltypes.SmallInteger(),
+        sql.sqltypes.String(),
+        sql.sqltypes.Text(),
+        sql.sqltypes.Time(),
+        sql.sqltypes.Unicode(),  # type: ignore[no-untyped-call]
+        sql.sqltypes.UnicodeText(),  # type: ignore[no-untyped-call]
+        sql.sqltypes.Uuid(),
+    ),
+)
+def test_sqlalchemy_custom_type_from_type_decorator(impl_: types.TypeEngine) -> None:
+    class CustomType(types.TypeDecorator):
+        impl = impl_
+
+    class Base(orm.DeclarativeBase):
+        type_annotation_map = {object: CustomType}
+
+    class Model(Base):
+        __tablename__ = "model_with_custom_types"
+
+        id: orm.Mapped[int] = orm.mapped_column(primary_key=True)
+        custom_type: orm.Mapped[Any] = orm.mapped_column(type_=CustomType(), nullable=False)
+        custom_type_from_annotation_map: orm.Mapped[object]
+
+    class ModelFactory(SQLAlchemyFactory[Model]):
+        __model__ = Model
+
+    instance = ModelFactory.build()
+    assert isinstance(instance.id, int)
+    assert isinstance(instance.custom_type, impl_.python_type)
+    assert isinstance(instance.custom_type_from_annotation_map, impl_.python_type)
+
+
+def test_sqlalchemy_custom_type_from_user_defined_type__overridden() -> None:
+    class CustomType(types.UserDefinedType):
+        ...
+
+    class Base(orm.DeclarativeBase):
+        ...
+
+    class Model(Base):
+        __tablename__ = "model_with_custom_types"
+
+        id: orm.Mapped[int] = orm.mapped_column(primary_key=True)
+        custom_type: orm.Mapped[Any] = orm.mapped_column(type_=CustomType())
+
+    class ModelFactory(SQLAlchemyFactory[Model]):
+        __model__ = Model
+
+        @classmethod
+        def get_sqlalchemy_types(cls) -> Dict[Any, Callable[[], Any]]:
+            return {**super().get_sqlalchemy_types(), CustomType: lambda: cls.__faker__.date_time()}
+
+    instance = ModelFactory.build()
+    assert isinstance(instance.id, int)
+    assert isinstance(instance.custom_type, datetime.datetime)
+
+
+def test_sqlalchemy_custom_type_from_user_defined_type__type_not_supported() -> None:
+    class CustomType(types.UserDefinedType):
+        ...
+
+    class Base(orm.DeclarativeBase):
+        ...
+
+    class Model(Base):
+        __tablename__ = "model_with_custom_types"
+
+        id: orm.Mapped[int] = orm.mapped_column(primary_key=True)
+        custom_type: orm.Mapped[Any] = orm.mapped_column(type_=CustomType())
+
+    class ModelFactory(SQLAlchemyFactory[Model]):
+        __model__ = Model
+
+    with pytest.raises(ParameterException, match="User defined type detected"):
+        ModelFactory.build()