From a99f7d81cd8bddfd4f089e79c4eef5c5bebccbd0 Mon Sep 17 00:00:00 2001 From: gazorby Date: Tue, 2 Jun 2026 20:04:15 +0200 Subject: [PATCH 1/3] fix: using strawberry.lazy(...) break circular refs --- src/strawchemy/utils/registry.py | 10 ++++-- src/strawchemy/utils/strawberry.py | 10 +++--- ...mas[forwardref_circular_default_scope].gql | 25 +++++++++++++++ ...emas[forwardref_circular_global_scope].gql | 28 ++++++++++++++++ ...y_schemas[lazy_circular_default_scope].gql | 25 +++++++++++++++ ...ry_schemas[lazy_circular_global_scope].gql | 28 ++++++++++++++++ tests/unit/mapping/test_schemas.py | 4 +++ tests/unit/schemas/forwardref/__init__.py | 0 tests/unit/schemas/forwardref/a.py | 32 +++++++++++++++++++ tests/unit/schemas/forwardref/b.py | 24 ++++++++++++++ tests/unit/schemas/forwardref/query.py | 12 +++++++ .../schemas/forwardref_global/__init__.py | 0 tests/unit/schemas/forwardref_global/a.py | 32 +++++++++++++++++++ tests/unit/schemas/forwardref_global/b.py | 20 ++++++++++++ tests/unit/schemas/forwardref_global/query.py | 12 +++++++ tests/unit/schemas/lazy/__init__.py | 0 tests/unit/schemas/lazy/a.py | 31 ++++++++++++++++++ tests/unit/schemas/lazy/b.py | 24 ++++++++++++++ tests/unit/schemas/lazy/query.py | 12 +++++++ tests/unit/schemas/lazy_global/__init__.py | 0 tests/unit/schemas/lazy_global/a.py | 32 +++++++++++++++++++ tests/unit/schemas/lazy_global/b.py | 20 ++++++++++++ tests/unit/schemas/lazy_global/query.py | 12 +++++++ 23 files changed, 386 insertions(+), 7 deletions(-) create mode 100644 tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[forwardref_circular_default_scope].gql create mode 100644 tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[forwardref_circular_global_scope].gql create mode 100644 tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[lazy_circular_default_scope].gql create mode 100644 tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[lazy_circular_global_scope].gql create mode 100644 tests/unit/schemas/forwardref/__init__.py create mode 100644 tests/unit/schemas/forwardref/a.py create mode 100644 tests/unit/schemas/forwardref/b.py create mode 100644 tests/unit/schemas/forwardref/query.py create mode 100644 tests/unit/schemas/forwardref_global/__init__.py create mode 100644 tests/unit/schemas/forwardref_global/a.py create mode 100644 tests/unit/schemas/forwardref_global/b.py create mode 100644 tests/unit/schemas/forwardref_global/query.py create mode 100644 tests/unit/schemas/lazy/__init__.py create mode 100644 tests/unit/schemas/lazy/a.py create mode 100644 tests/unit/schemas/lazy/b.py create mode 100644 tests/unit/schemas/lazy/query.py create mode 100644 tests/unit/schemas/lazy_global/__init__.py create mode 100644 tests/unit/schemas/lazy_global/a.py create mode 100644 tests/unit/schemas/lazy_global/b.py create mode 100644 tests/unit/schemas/lazy_global/query.py diff --git a/src/strawchemy/utils/registry.py b/src/strawchemy/utils/registry.py index ab5e590f..33494205 100644 --- a/src/strawchemy/utils/registry.py +++ b/src/strawchemy/utils/registry.py @@ -7,6 +7,7 @@ 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 @@ -177,8 +178,11 @@ 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 @@ -222,7 +226,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) diff --git a/src/strawchemy/utils/strawberry.py b/src/strawchemy/utils/strawberry.py index 75d52bae..6f13c538 100644 --- a/src/strawchemy/utils/strawberry.py +++ b/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_,) diff --git a/tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[forwardref_circular_default_scope].gql b/tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[forwardref_circular_default_scope].gql new file mode 100644 index 00000000..b791c348 --- /dev/null +++ b/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 +''' \ No newline at end of file diff --git a/tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[forwardref_circular_global_scope].gql b/tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[forwardref_circular_global_scope].gql new file mode 100644 index 00000000..a661e2ce --- /dev/null +++ b/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 +''' \ No newline at end of file diff --git a/tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[lazy_circular_default_scope].gql b/tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[lazy_circular_default_scope].gql new file mode 100644 index 00000000..b791c348 --- /dev/null +++ b/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 +''' \ No newline at end of file diff --git a/tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[lazy_circular_global_scope].gql b/tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[lazy_circular_global_scope].gql new file mode 100644 index 00000000..a661e2ce --- /dev/null +++ b/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 +''' \ No newline at end of file diff --git a/tests/unit/mapping/test_schemas.py b/tests/unit/mapping/test_schemas.py index 39479d55..1cadbb3b 100644 --- a/tests/unit/mapping/test_schemas.py +++ b/tests/unit/mapping/test_schemas.py @@ -300,6 +300,10 @@ 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.mark.snapshot diff --git a/tests/unit/schemas/forwardref/__init__.py b/tests/unit/schemas/forwardref/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/schemas/forwardref/a.py b/tests/unit/schemas/forwardref/a.py new file mode 100644 index 00000000..c46f91f7 --- /dev/null +++ b/tests/unit/schemas/forwardref/a.py @@ -0,0 +1,32 @@ +"""Circular plain forward-reference annotations between two `@strawchemy.type` classes (default scope). + +Same circular structure as `tests/unit/schemas/lazy/`, but using plain string forward +annotations (no `strawberry.lazy`). With `from __future__ import annotations` the +return annotation is a bare string `"BNode"` resolved by the registry's forward-ref +machinery at schema build. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import strawberry +from strawberry import auto + +from strawchemy import Strawchemy +from tests.unit.models import Color + +strawchemy = Strawchemy("postgresql") + +if TYPE_CHECKING: + from tests.unit.schemas.forwardref.b import BNode + + +@strawchemy.type(Color) +class ANode: + id: auto + name: auto + + @strawberry.field + def b(self) -> BNode: + raise NotImplementedError diff --git a/tests/unit/schemas/forwardref/b.py b/tests/unit/schemas/forwardref/b.py new file mode 100644 index 00000000..9affe493 --- /dev/null +++ b/tests/unit/schemas/forwardref/b.py @@ -0,0 +1,24 @@ +"""See `tests/unit/schemas/forwardref/a.py` — circular plain forward references, default scope.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import strawberry +from strawberry import auto + +from tests.unit.models import Fruit +from tests.unit.schemas.forwardref.a import strawchemy + +if TYPE_CHECKING: + from tests.unit.schemas.forwardref.a import ANode + + +@strawchemy.type(Fruit) +class BNode: + id: auto + name: auto + + @strawberry.field + def a(self) -> ANode: + raise NotImplementedError diff --git a/tests/unit/schemas/forwardref/query.py b/tests/unit/schemas/forwardref/query.py new file mode 100644 index 00000000..cf0a05cc --- /dev/null +++ b/tests/unit/schemas/forwardref/query.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +import strawberry + +from tests.unit.schemas.forwardref.a import ANode, strawchemy +from tests.unit.schemas.forwardref.b import BNode + + +@strawberry.type +class Query: + a: ANode = strawchemy.field() + b: BNode = strawchemy.field() diff --git a/tests/unit/schemas/forwardref_global/__init__.py b/tests/unit/schemas/forwardref_global/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/schemas/forwardref_global/a.py b/tests/unit/schemas/forwardref_global/a.py new file mode 100644 index 00000000..89bc97d7 --- /dev/null +++ b/tests/unit/schemas/forwardref_global/a.py @@ -0,0 +1,32 @@ +"""Circular plain forward-reference annotations with global-scope canonicalization. + +Same as `tests/unit/schemas/lazy_global/` but using plain string forward annotations +instead of `strawberry.lazy`. `ColorNode` is the `scope="schema"` canonical Color type; +`FruitNode` (include="all") auto-generates a `color` reference that canonicalizes to +`ColorNode` (no duplicate `ColorType`). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import strawberry +from strawberry import auto + +from strawchemy import Strawchemy +from tests.unit.models import Color + +strawchemy = Strawchemy("postgresql") + +if TYPE_CHECKING: + from tests.unit.schemas.forwardref_global.b import FruitNode + + +@strawchemy.type(Color, scope="schema") +class ColorNode: + id: auto + name: auto + + @strawberry.field + def featured_fruit(self) -> FruitNode: + raise NotImplementedError diff --git a/tests/unit/schemas/forwardref_global/b.py b/tests/unit/schemas/forwardref_global/b.py new file mode 100644 index 00000000..d531eb26 --- /dev/null +++ b/tests/unit/schemas/forwardref_global/b.py @@ -0,0 +1,20 @@ +"""See `tests/unit/schemas/forwardref_global/a.py` — circular plain forward references, global scope.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import strawberry + +from tests.unit.models import Fruit +from tests.unit.schemas.forwardref_global.a import strawchemy + +if TYPE_CHECKING: + from tests.unit.schemas.forwardref_global.a import ColorNode + + +@strawchemy.type(Fruit, include="all") +class FruitNode: + @strawberry.field + def primary_color(self) -> ColorNode: + raise NotImplementedError diff --git a/tests/unit/schemas/forwardref_global/query.py b/tests/unit/schemas/forwardref_global/query.py new file mode 100644 index 00000000..6554d942 --- /dev/null +++ b/tests/unit/schemas/forwardref_global/query.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +import strawberry + +from tests.unit.schemas.forwardref_global.a import ColorNode, strawchemy +from tests.unit.schemas.forwardref_global.b import FruitNode + + +@strawberry.type +class Query: + fruit: FruitNode = strawchemy.field() + color: ColorNode = strawchemy.field() diff --git a/tests/unit/schemas/lazy/__init__.py b/tests/unit/schemas/lazy/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/schemas/lazy/a.py b/tests/unit/schemas/lazy/a.py new file mode 100644 index 00000000..3816911c --- /dev/null +++ b/tests/unit/schemas/lazy/a.py @@ -0,0 +1,31 @@ +"""Circular `strawberry.lazy` references between two `@strawchemy.type` classes (default scope). + +`ANode` (over `Color`) lazily references `BNode` in module `b`, which lazily +references `ANode` back. Both lazy fields resolve to their target types when the +schema is built. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Annotated + +import strawberry +from strawberry import auto + +from strawchemy import Strawchemy +from tests.unit.models import Color + +strawchemy = Strawchemy("postgresql") + +if TYPE_CHECKING: + from tests.unit.schemas.lazy.b import BNode + + +@strawchemy.type(Color) +class ANode: + id: auto + name: auto + + @strawberry.field + def b(self) -> Annotated[BNode, strawberry.lazy("tests.unit.schemas.lazy.b")]: + raise NotImplementedError diff --git a/tests/unit/schemas/lazy/b.py b/tests/unit/schemas/lazy/b.py new file mode 100644 index 00000000..838fb48e --- /dev/null +++ b/tests/unit/schemas/lazy/b.py @@ -0,0 +1,24 @@ +"""See `tests/unit/schemas/lazy/a.py` — circular `strawberry.lazy` references, default scope.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Annotated + +import strawberry +from strawberry import auto + +from tests.unit.models import Fruit +from tests.unit.schemas.lazy.a import strawchemy + +if TYPE_CHECKING: + from tests.unit.schemas.lazy.a import ANode + + +@strawchemy.type(Fruit) +class BNode: + id: auto + name: auto + + @strawberry.field + def a(self) -> Annotated[ANode, strawberry.lazy("tests.unit.schemas.lazy.a")]: + raise NotImplementedError diff --git a/tests/unit/schemas/lazy/query.py b/tests/unit/schemas/lazy/query.py new file mode 100644 index 00000000..fb72ba18 --- /dev/null +++ b/tests/unit/schemas/lazy/query.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +import strawberry + +from tests.unit.schemas.lazy.a import ANode, strawchemy +from tests.unit.schemas.lazy.b import BNode + + +@strawberry.type +class Query: + a: ANode = strawchemy.field() + b: BNode = strawchemy.field() diff --git a/tests/unit/schemas/lazy_global/__init__.py b/tests/unit/schemas/lazy_global/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/schemas/lazy_global/a.py b/tests/unit/schemas/lazy_global/a.py new file mode 100644 index 00000000..29440c44 --- /dev/null +++ b/tests/unit/schemas/lazy_global/a.py @@ -0,0 +1,32 @@ +"""Circular `strawberry.lazy` references with global-scope canonicalization. + +`ColorNode` is the `scope="schema"` canonical type for `Color`. `FruitNode` +(include="all") auto-generates a `color` reference that canonicalizes to `ColorNode` +(no duplicate `ColorType` in the schema), and the two types reference each other via +circular `strawberry.lazy(...)` fields. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Annotated + +import strawberry +from strawberry import auto + +from strawchemy import Strawchemy +from tests.unit.models import Color + +strawchemy = Strawchemy("postgresql") + +if TYPE_CHECKING: + from tests.unit.schemas.lazy_global.b import FruitNode + + +@strawchemy.type(Color, scope="schema") +class ColorNode: + id: auto + name: auto + + @strawberry.field + def featured_fruit(self) -> Annotated[FruitNode, strawberry.lazy("tests.unit.schemas.lazy_global.b")]: + raise NotImplementedError diff --git a/tests/unit/schemas/lazy_global/b.py b/tests/unit/schemas/lazy_global/b.py new file mode 100644 index 00000000..775b5edd --- /dev/null +++ b/tests/unit/schemas/lazy_global/b.py @@ -0,0 +1,20 @@ +"""See `tests/unit/schemas/lazy_global/a.py` — circular `strawberry.lazy` references, global scope.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Annotated + +import strawberry + +from tests.unit.models import Fruit +from tests.unit.schemas.lazy_global.a import strawchemy + +if TYPE_CHECKING: + from tests.unit.schemas.lazy_global.a import ColorNode + + +@strawchemy.type(Fruit, include="all") +class FruitNode: + @strawberry.field + def primary_color(self) -> Annotated[ColorNode, strawberry.lazy("tests.unit.schemas.lazy_global.a")]: + raise NotImplementedError diff --git a/tests/unit/schemas/lazy_global/query.py b/tests/unit/schemas/lazy_global/query.py new file mode 100644 index 00000000..228fb9e9 --- /dev/null +++ b/tests/unit/schemas/lazy_global/query.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +import strawberry + +from tests.unit.schemas.lazy_global.a import ColorNode, strawchemy +from tests.unit.schemas.lazy_global.b import FruitNode + + +@strawberry.type +class Query: + fruit: FruitNode = strawchemy.field() + color: ColorNode = strawchemy.field() From bb6719e67aee4ee778f04cbd6ed8d6862e1db547 Mon Sep 17 00:00:00 2001 From: gazorby Date: Sat, 6 Jun 2026 23:43:36 +0200 Subject: [PATCH 2/3] feat(registry): support replacing types within StrawberryUnions during override --- src/strawchemy/utils/registry.py | 102 ++++++++++++------ ...est_query_schemas[union_override_lazy].gql | 25 +++++ ...st_query_schemas[union_override_plain].gql | 25 +++++ tests/unit/mapping/test_schemas.py | 2 + tests/unit/schemas/union_override_lazy.py | 46 ++++++++ tests/unit/schemas/union_override_plain.py | 45 ++++++++ 6 files changed, 213 insertions(+), 32 deletions(-) create mode 100644 tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[union_override_lazy].gql create mode 100644 tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[union_override_plain].gql create mode 100644 tests/unit/schemas/union_override_lazy.py create mode 100644 tests/unit/schemas/union_override_plain.py diff --git a/src/strawchemy/utils/registry.py b/src/strawchemy/utils/registry.py index 33494205..8b254c56 100644 --- a/src/strawchemy/utils/registry.py +++ b/src/strawchemy/utils/registry.py @@ -12,6 +12,7 @@ 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 @@ -51,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. + """Reference to a type used by a field or argument, used to update it later.""" - Args: - container: The container to replace the type in. - strawberry_type: The type to replace with. - - 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 + """ + 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. """ - 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: + 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: @@ -92,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) @@ -186,7 +224,7 @@ def _update_references(self, field: StrawberryField | StrawberryArgument, graphq 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: @@ -303,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, @@ -390,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, @@ -433,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, diff --git a/tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[union_override_lazy].gql b/tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[union_override_lazy].gql new file mode 100644 index 00000000..05351c1b --- /dev/null +++ b/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 +''' \ No newline at end of file diff --git a/tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[union_override_plain].gql b/tests/unit/mapping/__snapshots__/test_schemas/test_query_schemas[union_override_plain].gql new file mode 100644 index 00000000..05351c1b --- /dev/null +++ b/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 +''' \ No newline at end of file diff --git a/tests/unit/mapping/test_schemas.py b/tests/unit/mapping/test_schemas.py index 1cadbb3b..ae873030 100644 --- a/tests/unit/mapping/test_schemas.py +++ b/tests/unit/mapping/test_schemas.py @@ -304,6 +304,8 @@ def test_update_mutation_by_filter_type_not_list_fail() -> None: 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 diff --git a/tests/unit/schemas/union_override_lazy.py b/tests/unit/schemas/union_override_lazy.py new file mode 100644 index 00000000..2e570d6f --- /dev/null +++ b/tests/unit/schemas/union_override_lazy.py @@ -0,0 +1,46 @@ +"""Union field whose override-type members must not rewrite the field type. + +``ParentNode.parent`` is declared as ``ChildUnion | None``, a Strawberry union of +``GroupNode`` and a ``strawberry.lazy`` reference to ``TagNode``. Registering the +``override=True`` member types must keep the field typed as the union, swapping only +the matching member inside it. +""" + +from __future__ import annotations + +from typing import Annotated, Union + +import strawberry + +from strawchemy import Strawchemy +from tests.unit.models import Group, Tag, User + +strawchemy = Strawchemy("postgresql") + + +@strawchemy.type(Group, include=["id"], override=True) +class GroupNode: + pass + + +ChildUnion = Annotated[ + Union[GroupNode, Annotated["TagNode", strawberry.lazy("tests.unit.schemas.union_override_lazy")]], + strawberry.union("ChildUnion"), +] + + +@strawchemy.type(User, include=["id"], override=True) +class ParentNode: + @strawberry.field(graphql_type=ChildUnion | None) + def parent(self) -> object | None: + return None + + +@strawchemy.type(Tag, include=["id"], override=True) +class TagNode: + pass + + +@strawberry.type +class Query: + parent: ParentNode | None = None diff --git a/tests/unit/schemas/union_override_plain.py b/tests/unit/schemas/union_override_plain.py new file mode 100644 index 00000000..d8824988 --- /dev/null +++ b/tests/unit/schemas/union_override_plain.py @@ -0,0 +1,45 @@ +"""Union field whose override-type members must not rewrite the field type (no lazy). + +Same shape as ``union_override_lazy`` but both union members are direct class +references. Registering the ``override=True`` member types must keep the field typed +as the union. +""" + +from __future__ import annotations + +from typing import Annotated, Union + +import strawberry + +from strawchemy import Strawchemy +from tests.unit.models import Group, Tag, User + +strawchemy = Strawchemy("postgresql") + + +@strawchemy.type(Group, include=["id"], override=True) +class GroupNode: + pass + + +@strawchemy.type(Tag, include=["id"], override=True) +class TagNode: + pass + + +ChildUnion = Annotated[ + Union[GroupNode, TagNode], + strawberry.union("ChildUnion"), +] + + +@strawchemy.type(User, include=["id"], override=True) +class ParentNode: + @strawberry.field(graphql_type=ChildUnion | None) + def parent(self) -> object | None: + return None + + +@strawberry.type +class Query: + parent: ParentNode | None = None From fe0f8ad4d5b0ebb858526f986d6d10e5047edd08 Mon Sep 17 00:00:00 2001 From: gazorby Date: Tue, 9 Jun 2026 20:54:21 +0200 Subject: [PATCH 3/3] test(registry): add test_type_reference_matches_target case --- tests/unit/test_registry.py | 45 +++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 tests/unit/test_registry.py diff --git a/tests/unit/test_registry.py b/tests/unit/test_registry.py new file mode 100644 index 00000000..35ecbd38 --- /dev/null +++ b/tests/unit/test_registry.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from typing import Any + +import pytest +import strawberry +from strawberry import LazyType +from strawberry.types import get_object_definition + +from strawchemy.utils.registry import _TypeReference + + +@strawberry.type +class GroupNode: + id: int + + +@strawberry.type(name="TagNode") +class RenamedTag: + id: int + + +class NotStrawberry: + pass + + +_REF_HOLDER = get_object_definition(GroupNode, strict=True).fields[0] +_NOT_THE_MEMBER = object() + + +@pytest.mark.parametrize( + ("target", "target_name", "member", "expected"), + [ + pytest.param(GroupNode, "GroupNode", GroupNode, True, id="identity-match"), + pytest.param(_NOT_THE_MEMBER, None, GroupNode, False, id="no-target-name"), + pytest.param(_NOT_THE_MEMBER, "TagNode", LazyType("TagNode", "some.module"), True, id="lazy-name-match"), + pytest.param(_NOT_THE_MEMBER, "TagNode", LazyType("Other", "some.module"), False, id="lazy-name-mismatch"), + pytest.param(_NOT_THE_MEMBER, "TagNode", RenamedTag, True, id="definition-name-match"), + pytest.param(_NOT_THE_MEMBER, "TagNode", GroupNode, False, id="definition-name-mismatch"), + pytest.param(_NOT_THE_MEMBER, "NotStrawberry", NotStrawberry, False, id="no-object-definition"), + ], +) +def test_type_reference_matches_target(target: Any, target_name: str | None, member: Any, expected: bool) -> None: + reference = _TypeReference(ref_holder=_REF_HOLDER, target=target, target_name=target_name) + assert reference._matches_target(member) is expected # noqa: SLF001