fix: fix generics resolution when using Annotated unions (#3950)

* fix: fix generics resolution when using Annotated unions

Fix #3289

* fix mypy

Author: bellini666

RELEASE.md

@@ -0,0 +1,46 @@
+Release type: patch
+
+This release fixes the resolution of `Generics` when specializing using a union
+defined with `Annotated`, like in the example below:
+
+```python
+from typing import Annotated, Generic, TypeVar, Union
+import strawberry
+
+T = TypeVar("T")
+
+
+@strawberry.type
+class User:
+    name: str
+    age: int
+
+
+@strawberry.type
+class ProUser:
+    name: str
+    age: float
+
+
+@strawberry.type
+class GenType(Generic[T]):
+    data: T
+
+
+GeneralUser = Annotated[Union[User, ProUser], strawberry.union("GeneralUser")]
+
+
+@strawberry.type
+class Response(GenType[GeneralUser]): ...
+
+
+@strawberry.type
+class Query:
+    @strawberry.field
+    def user(self) -> Response: ...
+
+
+schema = strawberry.Schema(query=Query)
+```
+
+Before this would raise a `TypeError`, now it works as expected.

strawberry/annotation.py

@@ -125,16 +125,44 @@ def _get_type_with_args(
 
         return evaled_type, []
 
-    def resolve(self) -> Union[StrawberryType, type]:
+    def resolve(
+        self,
+        *,
+        type_definition: Optional[StrawberryObjectDefinition] = None,
+    ) -> Union[StrawberryType, type]:
         """Return resolved (transformed) annotation."""
-        if self.__resolve_cache__ is None:
-            self.__resolve_cache__ = self._resolve()
+        if (resolved := self.__resolve_cache__) is None:
+            resolved = self._resolve()
+            self.__resolve_cache__ = resolved
+
+        # If this is a generic field, try to resolve it using its origin's
+        # specialized type_var_map
+        if self._is_type_generic(resolved) and type_definition is not None:
+            from strawberry.types.base import StrawberryType
+
+            specialized_type_var_map = type_definition.specialized_type_var_map
+            if specialized_type_var_map and isinstance(resolved, StrawberryType):
+                resolved = resolved.copy_with(specialized_type_var_map)
+
+            # If the field is still generic, try to resolve it from the type_definition
+            # that is asking for it.
+            if (
+                self._is_type_generic(resolved)
+                and type_definition.type_var_map
+                and isinstance(resolved, StrawberryType)
+            ):
+                resolved = resolved.copy_with(type_definition.type_var_map)
+
+            # Resolve the type again to resolve any `Annotated` types
+            resolved = self._resolve_evaled_type(resolved)
 
-        return self.__resolve_cache__
+        return resolved
 
     def _resolve(self) -> Union[StrawberryType, type]:
         evaled_type = cast("Any", self.evaluate())
+        return self._resolve_evaled_type(evaled_type)
 
+    def _resolve_evaled_type(self, evaled_type: Any) -> Union[StrawberryType, type]:
         if is_private(evaled_type):
             return evaled_type
 
@@ -145,7 +173,7 @@ def _resolve(self) -> Union[StrawberryType, type]:
         if self._is_lazy_type(evaled_type):
             return evaled_type
         if self._is_streamable(evaled_type, args):
-            return self.create_list(list[evaled_type])  # type: ignore[valid-type]
+            return self.create_list(list[evaled_type])
         if self._is_list(evaled_type):
             return self.create_list(evaled_type)
         if self._is_maybe(evaled_type):
@@ -292,6 +320,20 @@ def _is_enum(cls, annotation: Any) -> bool:
             return False
         return issubclass(annotation, Enum)
 
+    @classmethod
+    def _is_type_generic(cls, type_: Union[StrawberryType, type]) -> bool:
+        """Returns True if `resolver_type` is generic else False."""
+        from strawberry.types.base import StrawberryType
+
+        if isinstance(type_, StrawberryType):
+            return type_.is_graphql_generic
+
+        # solves the Generic subclass case
+        if has_object_definition(type_):
+            return type_.__strawberry_definition__.is_graphql_generic
+
+        return False
+
     @classmethod
     def _is_graphql_generic(cls, annotation: Any) -> bool:
         if hasattr(annotation, "__origin__"):

strawberry/types/field.py

@@ -13,7 +13,6 @@
     Optional,
     TypeVar,
     Union,
-    cast,
     overload,
 )
 
@@ -348,34 +347,14 @@ def resolve_type(
         with contextlib.suppress(NameError):
             # Prioritise the field type over the resolver return type
             if self.type_annotation is not None:
-                resolved = self.type_annotation.resolve()
+                resolved = self.type_annotation.resolve(type_definition=type_definition)
             elif self.base_resolver is not None and self.base_resolver.type is not None:
                 # Handle unannotated functions (such as lambdas)
                 # Generics will raise MissingTypesForGenericError later
                 # on if we let it be returned. So use `type_annotation` instead
                 # which is the same behaviour as having no type information.
                 resolved = self.base_resolver.type
 
-        # If this is a generic field, try to resolve it using its origin's
-        # specialized type_var_map
-        # TODO: should we check arguments here too?
-        if _is_generic(resolved):  # type: ignore
-            specialized_type_var_map = (
-                type_definition and type_definition.specialized_type_var_map
-            )
-            if specialized_type_var_map and isinstance(resolved, StrawberryType):
-                resolved = resolved.copy_with(specialized_type_var_map)
-
-            # If the field is still generic, try to resolve it from the type_definition
-            # that is asking for it.
-            if (
-                _is_generic(cast("Union[StrawberryType, type]", resolved))
-                and type_definition is not None
-                and type_definition.type_var_map
-                and isinstance(resolved, StrawberryType)
-            ):
-                resolved = resolved.copy_with(type_definition.type_var_map)
-
         return resolved
 
     def copy_with(

tests/schema/test_union.py

@@ -1194,3 +1194,60 @@ class Query:
     """
         ).strip()
     )
+
+
+def test_union_used_inside_generic():
+    T = TypeVar("T")
+
+    @strawberry.type
+    class User:
+        name: str
+        age: int
+
+    @strawberry.type
+    class ProUser:
+        name: str
+        age: float
+
+    @strawberry.type
+    class GenType(Generic[T]):
+        data: T
+
+    GeneralUser = Annotated[Union[User, ProUser], strawberry.union("GeneralUser")]
+
+    @strawberry.type
+    class Response(GenType[GeneralUser]): ...
+
+    @strawberry.type
+    class Query:
+        @strawberry.field
+        def user(self) -> Response: ...
+
+    schema = strawberry.Schema(query=Query)
+
+    assert (
+        str(schema)
+        == textwrap.dedent(
+            """
+        union GeneralUser = User | ProUser
+
+        type ProUser {
+          name: String!
+          age: Float!
+        }
+
+        type Query {
+          user: Response!
+        }
+
+        type Response {
+          data: GeneralUser!
+        }
+
+        type User {
+          name: String!
+          age: Int!
+        }
+    """
+        ).strip()
+    )