fix: using strawberry.lazy(...) break circular refs

src/strawchemy/utils/registry.py

@@ -7,10 +7,12 @@
 from typing import TYPE_CHECKING, Any, ForwardRef, Literal, NewType, TypeVar, cast, overload
 
 import strawberry
+from strawberry import LazyType
 from strawberry.annotation import StrawberryAnnotation
 from strawberry.types import get_object_definition, has_object_definition
 from strawberry.types.base import StrawberryContainer
 from strawberry.types.field import StrawberryField
+from strawberry.types.union import StrawberryUnion
 
 from strawchemy.dto.strawberry import MappedStrawberryGraphQLDTO
 from strawchemy.dto.types import cast_include_fields, is_fields_iterable
@@ -50,32 +52,71 @@
 _RegistryMissing = NewType("_RegistryMissing", object)
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True, slots=True)
 class _TypeReference:
-    ref_holder: StrawberryField | StrawberryArgument
-
-    @classmethod
-    def _replace_contained_type(
-        cls, container: StrawberryContainer, strawberry_type: type[WithStrawberryObjectDefinition]
-    ) -> StrawberryContainer:
-        """Recursively replace the contained type in a StrawberryContainer.
-
-        Args:
-            container: The container to replace the type in.
-            strawberry_type: The type to replace with.
+    """Reference to a type used by a field or argument, used to update it later."""
 
-        Returns:
-            A new container with the type replaced.
+    ref_holder: StrawberryField | StrawberryArgument
+    """The field or argument whose type holds the referenced type."""
+    target: Any = None
+    """The inner type (class, `LazyType`, forward ref) this reference was created for.
+
+    Used to match the right member when the referenced type is part of a `StrawberryUnion`.
+    """
+    target_name: str | None = None
+    """The registry name the reference is keyed under.
+
+    Fallback for matching a union member when `target` identity does not hold
+    (e.g. the member annotation was re-resolved).
+    """
+
+    def _matches_target(self, member: Any) -> bool:
+        """Check whether a union member is the one this reference was created for."""
+        if member is self.target:
+            return True
+        if self.target_name is None:
+            return False
+        if isinstance(member, LazyType):
+            return member.type_name == self.target_name
+        member_definition = get_object_definition(member)
+        return member_definition is not None and member_definition.name == self.target_name
+
+    def _replaced_union(
+        self, union: StrawberryUnion, strawberry_type: type[WithStrawberryObjectDefinition]
+    ) -> StrawberryUnion:
+        """Return a copy of the union with only the matching member replaced.
+
+        The member is matched by identity against the reference target, falling back
+        to a name comparison. If no member matches, the union is returned unchanged
         """
-        container_copy = copy(container)
-        if isinstance(container.of_type, StrawberryContainer):
-            replaced = cls._replace_contained_type(container.of_type, strawberry_type)
-        else:
-            replaced = strawberry_type
-        container_copy.of_type = replaced
-        return container_copy
-
-    def _set_type(self, strawberry_type: type[WithStrawberryObjectDefinition] | StrawberryContainer) -> None:
+        annotations = list(union.type_annotations)
+        for index, member in enumerate(union.types):
+            if self._matches_target(member):
+                annotations[index] = StrawberryAnnotation(strawberry_type, namespace=annotations[index].namespace)
+                union_copy = copy(union)
+                union_copy.type_annotations = tuple(annotations)
+                return union_copy
+        return union
+
+    def _replaced(
+        self, node: Any, strawberry_type: type[WithStrawberryObjectDefinition]
+    ) -> type[WithStrawberryObjectDefinition] | StrawberryContainer | StrawberryUnion:
+        """Recursively replace the referenced type within containers and unions.
+
+        Containers are copied with their inner type replaced; unions are rebuilt with
+        only the matching member swapped; any other node is replaced directly.
+        """
+        if isinstance(node, StrawberryContainer):
+            container_copy = copy(node)
+            container_copy.of_type = self._replaced(node.of_type, strawberry_type)
+            return container_copy
+        if isinstance(node, StrawberryUnion):
+            return self._replaced_union(node, strawberry_type)
+        return strawberry_type
+
+    def _set_type(
+        self, strawberry_type: type[WithStrawberryObjectDefinition] | StrawberryContainer | StrawberryUnion
+    ) -> None:
         """Set the type of the referenced field or argument.
 
         Args:
@@ -91,15 +132,13 @@ def _set_type(self, strawberry_type: type[WithStrawberryObjectDefinition] | Stra
     def update_type(self, strawberry_type: type[WithStrawberryObjectDefinition]) -> None:
         """Update the type of the referenced field or argument.
 
-        If the referenced type is a container, it will recursively replace the contained type.
+        Containers are recursed into; union members are replaced individually (a
+        union field is never replaced wholesale).
 
         Args:
             strawberry_type: The type to update to.
         """
-        if isinstance(self.ref_holder.type, StrawberryContainer):
-            self._set_type(self._replace_contained_type(self.ref_holder.type, strawberry_type))
-        else:
-            self._set_type(strawberry_type)
+        self._set_type(self._replaced(self.ref_holder.type, strawberry_type))
 
 
 @dataclasses.dataclass(frozen=True, eq=True)
@@ -177,12 +216,15 @@ def _update_references(self, field: StrawberryField | StrawberryArgument, graphq
             field: The field or argument to update the references of.
             graphql_type: The graphql type of the field.
         """
-        for inner_type in strawberry_contained_types(field.type):
-            field_type_name = self._get_field_type_name(field, inner_type, graphql_type)
+        for inner_type in strawberry_contained_types(field.type, resolve_lazy=False):
+            if isinstance(inner_type, LazyType):
+                field_type_name: str | None = inner_type.type_name
+            else:
+                field_type_name = self._get_field_type_name(field, inner_type, graphql_type)
             if not field_type_name:
                 continue
 
-            type_ref = _TypeReference(field)
+            type_ref = _TypeReference(field, target=inner_type, target_name=field_type_name)
             type_info = self.get(graphql_type, field_type_name, None)
 
             if type_info and not type_info.exclude_from_scope:
@@ -222,7 +264,7 @@ def _track_references(
             for argument in field.arguments:
                 if any(
                     get_object_definition(inner_type) is not None
-                    for inner_type in strawberry_contained_types(argument.type)
+                    for inner_type in strawberry_contained_types(argument.type, resolve_lazy=False)
                 ):
                     self._update_references(argument, "input")
             self._update_references(field, graphql_type)
@@ -299,7 +341,7 @@ def _name_clash(self, type_info: RegistryTypeInfo) -> bool:
             and not type_info.override
         )
 
-    def _get_type_info(
+    def _type_info(
         self,
         dto: type[StrawchemyObject | Enum],
         graphql_type: GraphQLType,
@@ -386,7 +428,7 @@ def register_type(
         description: str | None = None,
         directives: Sequence[object] | None = (),
     ) -> type[StrawchemyDTOT]:
-        type_info = self._get_type_info(
+        type_info = self._type_info(
             dto=dto,
             graphql_type=graphql_type,
             dto_config=dto_config,
@@ -429,7 +471,7 @@ def register_enum(
         description: str | None = None,
         directives: Sequence[object] = (),
     ) -> type[EnumT]:
-        type_info = self._get_type_info(
+        type_info = self._type_info(
             dto=enum_type,
             graphql_type="enum",
             dto_config=dto_config,

src/strawchemy/utils/strawberry.py

@@ -44,15 +44,17 @@ def dto_model_from_type(type_: Any) -> Any:
     return type_.__dto_model__
 
 
-def strawberry_contained_types(type_: StrawberryType | Any) -> tuple[Any, ...]:
+def strawberry_contained_types(type_: StrawberryType | Any, resolve_lazy: bool = True) -> tuple[Any, ...]:
     if isinstance(type_, LazyType):
-        return strawberry_contained_types(type_.resolve_type())
+        if not resolve_lazy:
+            return (type_,)
+        return strawberry_contained_types(type_.resolve_type(), resolve_lazy=resolve_lazy)
     if isinstance(type_, StrawberryContainer):
-        return strawberry_contained_types(type_.of_type)
+        return strawberry_contained_types(type_.of_type, resolve_lazy=resolve_lazy)
     if isinstance(type_, StrawberryUnion):
         union_types = []
         for union_type in type_.types:
-            union_types.extend(strawberry_contained_types(union_type))
+            union_types.extend(strawberry_contained_types(union_type, resolve_lazy=resolve_lazy))
         return tuple(union_types)
     return (type_,)
 

tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[forwardref_circular_default_scope].gql

@@ -0,0 +1,25 @@
+'''
+"""GraphQL type"""
+type ANode {
+  name: String!
+  id: UUID!
+  b: BNode!
+}
+
+"""GraphQL type"""
+type BNode {
+  name: String!
+  id: UUID!
+  a: ANode!
+}
+
+type Query {
+  """Fetch object from the ANode collection by id"""
+  a(id: UUID!): ANode!
+
+  """Fetch object from the BNode collection by id"""
+  b(id: UUID!): BNode!
+}
+
+scalar UUID
+'''

tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[forwardref_circular_global_scope].gql

@@ -0,0 +1,28 @@
+'''
+"""GraphQL type"""
+type ColorNode {
+  name: String!
+  id: UUID!
+  featuredFruit: FruitNode!
+}
+
+"""GraphQL type"""
+type FruitNode {
+  color: ColorNode!
+  name: String!
+  colorId: UUID
+  sweetness: Int!
+  id: UUID!
+  primaryColor: ColorNode!
+}
+
+type Query {
+  """Fetch object from the FruitNode collection by id"""
+  fruit(id: UUID!): FruitNode!
+
+  """Fetch object from the ColorNode collection by id"""
+  color(id: UUID!): ColorNode!
+}
+
+scalar UUID
+'''

tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[lazy_circular_default_scope].gql

@@ -0,0 +1,25 @@
+'''
+"""GraphQL type"""
+type ANode {
+  name: String!
+  id: UUID!
+  b: BNode!
+}
+
+"""GraphQL type"""
+type BNode {
+  name: String!
+  id: UUID!
+  a: ANode!
+}
+
+type Query {
+  """Fetch object from the ANode collection by id"""
+  a(id: UUID!): ANode!
+
+  """Fetch object from the BNode collection by id"""
+  b(id: UUID!): BNode!
+}
+
+scalar UUID
+'''

tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[lazy_circular_global_scope].gql

@@ -0,0 +1,28 @@
+'''
+"""GraphQL type"""
+type ColorNode {
+  name: String!
+  id: UUID!
+  featuredFruit: FruitNode!
+}
+
+"""GraphQL type"""
+type FruitNode {
+  color: ColorNode!
+  name: String!
+  colorId: UUID
+  sweetness: Int!
+  id: UUID!
+  primaryColor: ColorNode!
+}
+
+type Query {
+  """Fetch object from the FruitNode collection by id"""
+  fruit(id: UUID!): FruitNode!
+
+  """Fetch object from the ColorNode collection by id"""
+  color(id: UUID!): ColorNode!
+}
+
+scalar UUID
+'''

tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[union_override_lazy].gql

@@ -0,0 +1,25 @@
+'''
+union ChildUnion = GroupNode | TagNode
+
+"""GraphQL type"""
+type GroupNode {
+  id: UUID!
+}
+
+"""GraphQL type"""
+type ParentNode {
+  id: UUID!
+  parent: ChildUnion
+}
+
+type Query {
+  parent: ParentNode
+}
+
+"""GraphQL type"""
+type TagNode {
+  id: UUID!
+}
+
+scalar UUID
+'''

tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[union_override_plain].gql

@@ -0,0 +1,25 @@
+'''
+union ChildUnion = GroupNode | TagNode
+
+"""GraphQL type"""
+type GroupNode {
+  id: UUID!
+}
+
+"""GraphQL type"""
+type ParentNode {
+  id: UUID!
+  parent: ChildUnion
+}
+
+type Query {
+  parent: ParentNode
+}
+
+"""GraphQL type"""
+type TagNode {
+  id: UUID!
+}
+
+scalar UUID
+'''

tests/unit/mapping/test_schemas.py

@@ -300,6 +300,12 @@ def test_update_mutation_by_filter_type_not_list_fail() -> None:
         pytest.param("scope.schema_before.Query", id="scope_schema_before"),
         pytest.param("scope.schema_after.Query", id="scope_schema_after"),
         pytest.param("scope.schema_middle.Query", id="scope_schema_middle"),
+        pytest.param("lazy.query.Query", id="lazy_circular_default_scope"),
+        pytest.param("lazy_global.query.Query", id="lazy_circular_global_scope"),
+        pytest.param("forwardref.query.Query", id="forwardref_circular_default_scope"),
+        pytest.param("forwardref_global.query.Query", id="forwardref_circular_global_scope"),
+        pytest.param("union_override_lazy.Query", id="union_override_lazy"),
+        pytest.param("union_override_plain.Query", id="union_override_plain"),
     ],
 )
 @pytest.mark.snapshot