feat: support for custom collection_class in SQLAlchemy relationships (#776)

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

Author: AlexPetul

polyfactory/factories/sqlalchemy_factory.py

@@ -1,9 +1,22 @@
 from __future__ import annotations
 
+from collections.abc import Collection, Mapping
 from datetime import date, datetime
-from typing import TYPE_CHECKING, Annotated, Any, Callable, ClassVar, Generic, Protocol, TypeVar, Union
-
-from polyfactory.exceptions import MissingDependencyException, ParameterException
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Callable,
+    ClassVar,
+    Generic,
+    Protocol,
+    TypeVar,
+    Union,
+)
+
+from sqlalchemy.util.langhelpers import duck_type_collection
+
+from polyfactory.exceptions import ConfigurationException, MissingDependencyException, ParameterException
 from polyfactory.factories.base import BaseFactory
 from polyfactory.field_meta import Constraints, FieldMeta
 from polyfactory.persistence import AsyncPersistenceProtocol, SyncPersistenceProtocol
@@ -203,6 +216,30 @@ def get_type_from_column(cls, column: Column) -> type:
 
         return annotation
 
+    @classmethod
+    def get_type_from_collection_class(
+        cls,
+        collection_class: type[Collection[Any]] | Callable[[], Collection[Any]],
+        entity_class: Any,
+    ) -> type[Any]:
+        annotation: type[Any]
+
+        if isinstance(collection_class, type):
+            if issubclass(collection_class, Mapping):
+                annotation = dict[Any, entity_class]
+            else:
+                if not (duck_typed_as := duck_type_collection(collection_class)):
+                    msg = f"Cannot infer type from collection_class {collection_class}"
+                    raise ConfigurationException(
+                        msg,
+                    )
+
+                annotation = duck_typed_as[entity_class]  # pyright: ignore[reportIndexIssue]
+        else:
+            annotation = dict[Any, entity_class]
+
+        return annotation
+
     @classmethod
     def get_model_fields(cls) -> list[FieldMeta]:
         fields_meta: list[FieldMeta] = []
@@ -219,12 +256,22 @@ def get_model_fields(cls) -> list[FieldMeta]:
         if cls.__set_relationships__:
             for name, relationship in table.relationships.items():
                 class_ = relationship.entity.class_
-                annotation = class_ if not relationship.uselist else list[class_]  # type: ignore[valid-type]
+                annotation: Any
+
+                if relationship.uselist:
+                    collection_class = relationship.collection_class
+                    if collection_class is None:
+                        annotation = list[class_]  # type: ignore[valid-type]
+                    else:
+                        annotation = cls.get_type_from_collection_class(collection_class, class_)
+                else:
+                    annotation = class_
+
                 fields_meta.append(
                     FieldMeta.from_type(
                         name=name,
                         annotation=annotation,
-                    ),
+                    )
                 )
         if cls.__set_association_proxy__:
             for name, attr in table.all_orm_descriptors.items():

tests/sqlalchemy_factory/models.py

@@ -1,17 +1,7 @@
 from dataclasses import dataclass
 from typing import Any, Optional
 
-from sqlalchemy import (
-    Boolean,
-    Column,
-    DateTime,
-    ForeignKey,
-    Integer,
-    String,
-    func,
-    orm,
-    text,
-)
+from sqlalchemy import Boolean, Column, DateTime, ForeignKey, Integer, String, func, orm, text
 from sqlalchemy.ext.associationproxy import association_proxy
 from sqlalchemy.orm import relationship
 from sqlalchemy.orm.decl_api import DeclarativeMeta, registry
@@ -32,6 +22,20 @@ class Base(metaclass=DeclarativeMeta):
     metadata = _registry.metadata
 
 
+class CollectionParentMixin:
+    __abstract__ = True
+    __allow_unmapped__ = True
+
+    id = Column(Integer(), primary_key=True)
+
+
+class CollectionChildMixin:
+    __abstract__ = True
+    __allow_unmapped__ = True
+
+    id = Column(Integer(), primary_key=True)
+
+
 class Author(Base):
     __tablename__ = "authors"
 

tests/sqlalchemy_factory/test_sqlalchemy_factory_common.py

@@ -1,13 +1,15 @@
 import warnings
+from collections.abc import Collection
 from datetime import datetime
 from decimal import Decimal
 from enum import Enum
 from typing import Any, Callable, get_args
-from uuid import UUID
+from uuid import UUID, uuid4
 
 import pytest
 from sqlalchemy import (
     Column,
+    ForeignKey,
     Integer,
     Numeric,
     String,
@@ -20,6 +22,20 @@
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import Session
+
+try:
+    from sqlalchemy.orm.collections import (
+        attribute_keyed_dict,
+        column_keyed_dict,
+        keyfunc_mapping,
+    )
+except ImportError:  # SQLAlchemy < 2.0
+    from sqlalchemy.orm import collections as _collections
+
+    attribute_keyed_dict = _collections.attribute_mapped_collection  # type: ignore[attr-defined]
+    column_keyed_dict = _collections.column_mapped_collection  # type: ignore[attr-defined]
+    keyfunc_mapping = _collections.mapped_collection  # type: ignore[attr-defined]
+
 from sqlalchemy.orm.decl_api import DeclarativeMeta, registry
 
 from polyfactory.exceptions import ConfigurationException, ParameterException
@@ -32,9 +48,12 @@
     Author,
     Base,
     Book,
+    CollectionChildMixin,
+    CollectionParentMixin,
     NonSQLAchemyClass,
     _registry,
 )
+from tests.sqlalchemy_factory.types import ListLike, SetLike
 
 
 @pytest.mark.parametrize(
@@ -195,6 +214,184 @@ class AuthorFactory(SQLAlchemyFactory[Author]):
     assert isinstance(result.books[0], Book)
 
 
+@pytest.mark.parametrize(
+    "collection_class_type",
+    (set, list, ListLike, SetLike),
+)
+def test_relationship_collection_class_sequence(collection_class_type: type[Collection]) -> None:
+    table_suffix = uuid4().hex
+    _registry = registry()
+
+    class Base(metaclass=DeclarativeMeta):
+        __abstract__ = True
+        __allow_unmapped__ = True
+
+        registry = _registry
+        metadata = _registry.metadata
+
+    class Parent(CollectionParentMixin, Base):
+        __tablename__ = f"parent_{table_suffix}"
+
+        children: Any = orm.relationship("Child", collection_class=collection_class_type)
+
+    class Child(CollectionChildMixin, Base):
+        __tablename__ = f"child_{table_suffix}"
+
+        parent_id = Column(Integer(), ForeignKey(f"{Parent.__tablename__}.id"), nullable=False)
+
+    class ParentFactory(SQLAlchemyFactory[Parent]):
+        __model__ = Parent
+        __set_relationships__ = True
+
+    result = ParentFactory.build()
+    assert result.children is not None
+
+    assert isinstance(result.children, collection_class_type)
+    first_item = next(iter(result.children))
+    assert isinstance(first_item, Child)
+
+
+def test_relationship_collection_class_attribute_keyed_dict() -> None:
+    table_suffix = uuid4().hex
+    _registry = registry()
+
+    class Base(metaclass=DeclarativeMeta):
+        __abstract__ = True
+        __allow_unmapped__ = True
+
+        registry = _registry
+        metadata = _registry.metadata
+
+    class Parent(CollectionParentMixin, Base):
+        __tablename__ = f"parent_{table_suffix}"
+
+        children = orm.relationship("Child", collection_class=attribute_keyed_dict("id"))
+
+    class Child(CollectionChildMixin, Base):
+        __tablename__ = f"child_{table_suffix}"
+
+        parent_id = Column(Integer(), ForeignKey(f"{Parent.__tablename__}.id"), nullable=False)
+
+    class ParentFactory(SQLAlchemyFactory[Parent]):
+        __model__ = Parent
+        __set_relationships__ = True
+
+    result = ParentFactory.build()
+    assert result.children is not None
+
+    assert isinstance(result.children, dict)
+    child = next(iter(result.children.values()))
+    assert isinstance(child, Child)
+    assert child.id in result.children
+
+
+def test_relationship_collection_class_column_keyed_dict() -> None:
+    table_suffix = uuid4().hex
+    _registry = registry()
+
+    class Base(metaclass=DeclarativeMeta):
+        __abstract__ = True
+        __allow_unmapped__ = True
+
+        registry = _registry
+        metadata = _registry.metadata
+
+    class Parent(CollectionParentMixin, Base):
+        __tablename__ = f"parent_{table_suffix}"
+
+        children: Any
+
+    class Child(CollectionChildMixin, Base):
+        __tablename__ = f"child_{table_suffix}"
+
+        parent_id = Column(Integer(), ForeignKey(f"{Parent.__tablename__}.id"), nullable=False)
+
+    Parent.children = orm.relationship("Child", collection_class=column_keyed_dict(Child.__table__.c.id))  # type: ignore[attr-defined]
+
+    class ParentFactory(SQLAlchemyFactory[Parent]):
+        __model__ = Parent
+        __set_relationships__ = True
+
+    result = ParentFactory.build()
+    assert result.children is not None
+
+    assert isinstance(result.children, dict)
+    child = next(iter(result.children.values()))
+    assert isinstance(child, Child)
+    assert child.id in result.children
+
+
+def test_relationship_collection_class_arbitrary_keying() -> None:
+    table_suffix = uuid4().hex
+    _registry = registry()
+
+    class Base(metaclass=DeclarativeMeta):
+        __abstract__ = True
+        __allow_unmapped__ = True
+
+        registry = _registry
+        metadata = _registry.metadata
+
+    class Parent(CollectionParentMixin, Base):
+        __tablename__ = f"parent_{table_suffix}"
+
+        children = orm.relationship("Child", collection_class=keyfunc_mapping(lambda c: c.id))
+
+    class Child(CollectionChildMixin, Base):
+        __tablename__ = f"child_{table_suffix}"
+
+        parent_id = Column(Integer(), ForeignKey(f"{Parent.__tablename__}.id"), nullable=False)
+
+    class ParentFactory(SQLAlchemyFactory[Parent]):
+        __model__ = Parent
+        __set_relationships__ = True
+
+    result = ParentFactory.build()
+    assert result.children is not None
+
+    assert isinstance(result.children, dict)
+    child = next(iter(result.children.values()))
+    assert isinstance(child, Child)
+    assert child.id in result.children
+
+
+def test_relationship_collection_class_arbitrary_keyfunc() -> None:
+    def make_mapping() -> Any:
+        return keyfunc_mapping(lambda c: c.id)()  # type: ignore[call-arg]
+
+    table_suffix = uuid4().hex
+    _registry = registry()
+
+    class Base(metaclass=DeclarativeMeta):
+        __abstract__ = True
+        __allow_unmapped__ = True
+
+        registry = _registry
+        metadata = _registry.metadata
+
+    class Parent(CollectionParentMixin, Base):
+        __tablename__ = f"parent_{table_suffix}"
+
+        children = orm.relationship("Child", collection_class=make_mapping)
+
+    class Child(CollectionChildMixin, Base):
+        __tablename__ = f"child_{table_suffix}"
+
+        parent_id = Column(Integer(), ForeignKey(f"{Parent.__tablename__}.id"), nullable=False)
+
+    class ParentFactory(SQLAlchemyFactory[Parent]):
+        __model__ = Parent
+        __set_relationships__ = True
+
+    result = ParentFactory.build()
+    assert result.children is not None
+
+    assert isinstance(result.children, dict)
+    child = next(iter(result.children.values()))
+    assert isinstance(child, Child)
+    assert child.id in result.children
+
+
 def test_sqla_factory_create(engine: Engine) -> None:
     Base.metadata.create_all(engine)