Repository: litestar-org/polyfactory
Pull request #530 (2024-04-30 03:25:06)

feat(sqlfactory): support nested type in pg.array types and others
enhance sqlalchemy factory support postgresql dialects.
## Description
- support postgresql dialects types: `JSON`, `JSONB` and `HSTORE` type
- support nested pg.ARRAY types, such as `ARRAY(INET), ARRAY(MACADDR)` and etc
## Closes


wangxin688: seems `main` was broken for mypy. 
```
tests/test_recursive_models.py:56: error: Non-overlapping identity check (left operand type: "PydanticNode", right operand type: "type[_Sentinel]")  [comparison-overlap]
```
github-actions[bot]: Documentation preview will be available shortly at https://litestar-org.github.io/polyfactory-docs-preview/530
github-actions[bot]: Documentation preview will be available shortly at https://litestar-org.github.io/polyfactory-docs-preview/530
wangxin688: seems there's no way to reflect `JSON` or `JSONB` 's python type. so should I remove JSON, JSONB in types maping? or treat it as dict by default?. becuase JSON and JSONB also support `list[Any]`? what's you guys ideas?
peterschutt: > seems there's no way to reflect `JSON` or `JSONB` 's python type. so should I remove JSON, JSONB in types maping? or treat it as dict by default?. becuase JSON and JSONB also support `list[Any]`? what's you guys ideas?

Yeah, array is valid json - I'd have to defer to @guacs on what is the right way to handle this for polyfactory.
wangxin688: > > seems there's no way to reflect `JSON` or `JSONB` 's python type. so should I remove JSON, JSONB in types maping? or treat it as dict by default?. becuase JSON and JSONB also support `list[Any]`? what's you guys ideas?
> 
> Yeah, array is valid json - I'd have to defer to @guacs on what is the right way to handle this for polyfactory.

Just my ideas, maybe not the best way.
add a migac method to sqlafactory, such as `__default_json_type__: ClassVar[type] = dict`, but it can be change as other type when inherit from factory. then handle the `JSON` type by magic method definition.
github-actions[bot]: Documentation preview will be available shortly at https://litestar-org.github.io/polyfactory-docs-preview/530
github-actions[bot]: Documentation preview will be available shortly at https://litestar-org.github.io/polyfactory-docs-preview/530
wangxin688: > > > seems there's no way to reflect `JSON` or `JSONB` 's python type. so should I remove JSON, JSONB in types maping? or treat it as dict by default?. becuase JSON and JSONB also support `list[Any]`? what's you guys ideas?
> > 
> > 
> > Yeah, array is valid json - I'd have to defer to @guacs on what is the right way to handle this for polyfactory.
> 
> Just my ideas, maybe not the best way. add a migac method to sqlafactory, such as `__default_json_type__: ClassVar[type] = dict`, but it can be change as other type when inherit from factory. then handle the `JSON` type by magic method definition.

I have remove `JSON` and `JSONB` types due to uncertain python types. 
github-actions[bot]: Documentation preview will be available shortly at https://litestar-org.github.io/polyfactory-docs-preview/530
wangxin688: BTW, I didnot see `types.Uuid` in `get_sqlalchemy_types`, after my local test, it's well suited with `types.Uuid: uuid4` add to return dict.
guacs: @wangxin688 thanks for this!

For the JSON and JSONB, I think we can default to a random dictionary. If the users need something other than the deafult, they can override the `get_sqlalchemy_types` or `get_provider_map.`
wangxin688: looks like main branch is broken for lint again.
@guacs JSON has been added back. pls review
wangxin688: > looks like main branch is broken for lint again. @guacs JSON has been added back. pls review

may I ask why lint rule is checked for all files not newly changed files? 
wangxin688: > looks like main branch is broken for lint again. @guacs JSON has been added back. pls review

I executed `pre-commit run --all-files` on my PC, no issues were reported. As it's not related for this PR, how to skip linting checking? --
guacs: > > looks like main branch is broken for lint again. @guacs JSON has been added back. pls review
> 
> I executed `pre-commit run --all-files` on my PC, no issues were reported. As it's not related for this PR, how to skip linting checking? --

Skipping it isn't possible though it can be merged even if the linting fails. I'll fix the linting issues and try to make it only lint the changes files (sometime this week) and once that's merged, we can merge that into this branch?

Patch:
diff --git a/polyfactory/factories/sqlalchemy_factory.py b/polyfactory/factories/sqlalchemy_factory.py
index 5a2155b..4b62c8f 100644
--- a/polyfactory/factories/sqlalchemy_factory.py
+++ b/polyfactory/factories/sqlalchemy_factory.py
@@ -95,6 +95,9 @@ class SQLAlchemyFactory(Generic[T], BaseFactory[T]):
             postgresql.NUMRANGE: lambda: tuple(sorted([cls.__faker__.pyint(), cls.__faker__.pyint()])),
             postgresql.TSRANGE: lambda: (cls.__faker__.past_datetime(), datetime.now()),  # noqa: DTZ005
             postgresql.TSTZRANGE: lambda: (cls.__faker__.past_datetime(), datetime.now()),  # noqa: DTZ005
+            postgresql.HSTORE: lambda: cls.__faker__.pydict(),
+            # `types.JSON` is compatible for sqlachemy extend dialects. Such as `pg.JSON` and `JSONB`
+            types.JSON: lambda: cls.__faker__.pydict(),
         }
 
     @classmethod
@@ -124,8 +127,14 @@ class SQLAlchemyFactory(Generic[T], BaseFactory[T]):
     @classmethod
     def get_type_from_column(cls, column: Column) -> type:
         column_type = type(column.type)
-        if column_type in cls.get_sqlalchemy_types():
+        sqla_types = cls.get_sqlalchemy_types()
+        if column_type in sqla_types:
             annotation = column_type
+        elif issubclass(column_type, postgresql.ARRAY):
+            if type(column.type.item_type) in sqla_types:  # type: ignore[attr-defined]
+                annotation = List[type(column.type.item_type)]  # type: ignore[attr-defined,misc]
+            else:
+                annotation = List[column.type.item_type.python_type]  # type: ignore[assignment,name-defined]
         elif issubclass(column_type, types.ARRAY):
             annotation = List[column.type.item_type.python_type]  # type: ignore[assignment,name-defined]
         else:


Test patch:
diff --git a/tests/sqlalchemy_factory/test_sqlalchemy_factory_v2.py b/tests/sqlalchemy_factory/test_sqlalchemy_factory_v2.py
index ed293e7..de2ab7e 100644
--- a/tests/sqlalchemy_factory/test_sqlalchemy_factory_v2.py
+++ b/tests/sqlalchemy_factory/test_sqlalchemy_factory_v2.py
@@ -1,8 +1,15 @@
 from enum import Enum
-from typing import Any, List
+from ipaddress import ip_network
+from typing import Any, Dict, List
+from uuid import UUID
 
 import pytest
-from sqlalchemy import ForeignKey, __version__, orm, types
+from sqlalchemy import ForeignKey, Text, __version__, orm, types
+from sqlalchemy.dialects.mssql import JSON as MSSQL_JSON
+from sqlalchemy.dialects.mysql import JSON as MYSQL_JSON
+from sqlalchemy.dialects.postgresql import ARRAY, CIDR, HSTORE, INET, JSON, JSONB
+from sqlalchemy.dialects.sqlite import JSON as SQLITE_JSON
+from sqlalchemy.ext.mutable import MutableDict, MutableList
 
 from polyfactory.factories.sqlalchemy_factory import SQLAlchemyFactory
 
@@ -64,6 +71,63 @@ def test_python_type_handling_v2() -> None:
     assert isinstance(instance.str_array_type[0], str)
 
 
+def test_pg_dialect_types() -> None:
+    class Base(orm.DeclarativeBase): ...
+
+    class SqlaModel(Base):
+        __tablename__ = "sql_models"
+        id: orm.Mapped[int] = orm.mapped_column(primary_key=True)
+        uuid_type: orm.Mapped[UUID] = orm.mapped_column(type_=types.UUID)
+        nested_array_inet: orm.Mapped[List[str]] = orm.mapped_column(type_=ARRAY(INET, dimensions=1))
+        nested_array_cidr: orm.Mapped[List[str]] = orm.mapped_column(type_=ARRAY(CIDR, dimensions=1))
+        hstore_type: orm.Mapped[Dict] = orm.mapped_column(type_=HSTORE)
+        mut_nested_arry_inet: orm.Mapped[List[str]] = orm.mapped_column(
+            type_=MutableList.as_mutable(ARRAY(INET, dimensions=1))
+        )
+        pg_json_type: orm.Mapped[Dict] = orm.mapped_column(type_=JSON)
+        pg_jsonb_type: orm.Mapped[Dict] = orm.mapped_column(type_=JSONB)
+        common_json_type: orm.Mapped[Dict] = orm.mapped_column(type_=types.JSON)
+        mysql_json: orm.Mapped[Dict] = orm.mapped_column(type_=MYSQL_JSON)
+        sqlite_json: orm.Mapped[Dict] = orm.mapped_column(type_=SQLITE_JSON)
+        mssql_json: orm.Mapped[Dict] = orm.mapped_column(type_=MSSQL_JSON)
+
+        multible_pg_json_type: orm.Mapped[Dict] = orm.mapped_column(
+            type_=MutableDict.as_mutable(JSON(astext_type=Text()))  # type: ignore[no-untyped-call]
+        )
+        multible_pg_jsonb_type: orm.Mapped[Dict] = orm.mapped_column(
+            type_=MutableDict.as_mutable(JSONB(astext_type=Text()))  # type: ignore[no-untyped-call]
+        )
+        multible_common_json_type: orm.Mapped[Dict] = orm.mapped_column(type_=MutableDict.as_mutable(types.JSON()))
+        multible_mysql_json: orm.Mapped[Dict] = orm.mapped_column(type_=MutableDict.as_mutable(MYSQL_JSON()))
+        multible_sqlite_json: orm.Mapped[Dict] = orm.mapped_column(type_=MutableDict.as_mutable(SQLITE_JSON()))
+        multible_mssql_json: orm.Mapped[Dict] = orm.mapped_column(type_=MutableDict.as_mutable(MSSQL_JSON()))
+
+    class ModelFactory(SQLAlchemyFactory[SqlaModel]):
+        __model__ = SqlaModel
+
+    instance = ModelFactory.build()
+    assert isinstance(instance.nested_array_inet[0], str)
+    assert ip_network(instance.nested_array_inet[0])
+    assert isinstance(instance.nested_array_cidr[0], str)
+    assert ip_network(instance.nested_array_cidr[0])
+    assert isinstance(instance.hstore_type, dict)
+    assert isinstance(instance.uuid_type, UUID)
+    assert isinstance(instance.mut_nested_arry_inet[0], str)
+    assert ip_network(instance.mut_nested_arry_inet[0])
+    assert isinstance(instance.pg_json_type, dict)
+    assert isinstance(instance.pg_jsonb_type, dict)
+    assert isinstance(instance.common_json_type, dict)
+    assert isinstance(instance.mysql_json, dict)
+    assert isinstance(instance.sqlite_json, dict)
+    assert isinstance(instance.mssql_json, dict)
+    assert isinstance(instance.multible_pg_json_type, dict)
+    assert isinstance(instance.multible_pg_jsonb_type, dict)
+    assert isinstance(instance.multible_common_json_type, dict)
+    assert isinstance(instance.multible_mysql_json, dict)
+    assert isinstance(instance.multible_sqlite_json, dict)
+    assert isinstance(instance.multible_mssql_json, dict)
+
+
 @pytest.mark.parametrize(
     "type_",
     tuple(SQLAlchemyFactory.get_sqlalchemy_types().keys()),