Repository: strawberry-graphql/strawberry

Issue:
Add mypy plugin for strawberry.union
We have a function to create union types and also add name and description that are used when generating the GraphQL schema[1].

Unfortunately MyPy complains when using the annotation, since it doesn't recognise it as a type, here's an example:

```python
ExampleUnion = strawberry.union(
    "ExampleUnion",
    types=(
        ABC,
        CDE,
    ),
)

@strawberry.mutation
def abc(phone_number: str) -> ExampleUnion:
    # ...
```

Error:

```
Variable "api.mutations.abc.ExampleUnion" is not valid as a type
```

[1] This is why we can't use typing.Union directly, as we can't pass arguments to it


Discussion:

Resolving patch:
diff --git a/strawberry/ext/mypy_plugin.py b/strawberry/ext/mypy_plugin.py
--- a/strawberry/ext/mypy_plugin.py
+++ b/strawberry/ext/mypy_plugin.py
@@ -1,8 +1,23 @@
 from typing import Callable, Optional
 
-from mypy.plugin import AnalyzeTypeContext, ClassDefContext, Plugin
+from mypy.nodes import (
+    GDEF,
+    Expression,
+    IndexExpr,
+    NameExpr,
+    SymbolTableNode,
+    TupleExpr,
+    TypeAlias,
+)
+from mypy.plugin import (
+    AnalyzeTypeContext,
+    ClassDefContext,
+    DynamicClassDefContext,
+    Plugin,
+    SemanticAnalyzerPluginInterface,
+)
 from mypy.plugins import dataclasses
-from mypy.types import Type
+from mypy.types import Type, UnionType
 
 
 def lazy_type_analyze_callback(ctx: AnalyzeTypeContext) -> Type:
@@ -19,7 +34,48 @@ def private_type_analyze_callback(ctx: AnalyzeTypeContext) -> Type:
     return type_
 
 
+def _get_type_for_expr(expr: Expression, api: SemanticAnalyzerPluginInterface):
+    if isinstance(expr, NameExpr):
+        return api.named_type(expr.name)
+
+    if isinstance(expr, IndexExpr):
+        type_ = _get_type_for_expr(expr.base, api)
+        type_.args = [_get_type_for_expr(expr.index, api)]
+
+        return type_
+
+    raise ValueError(f"Unsupported expression f{type(expr)}")
+
+
+def union_hook(ctx: DynamicClassDefContext) -> None:
+    types = ctx.call.args[1]
+
+    if isinstance(types, TupleExpr):
+        type_ = UnionType(tuple(_get_type_for_expr(x, ctx.api) for x in types.items))
+
+        type_alias = TypeAlias(
+            type_,
+            fullname=ctx.api.qualified_name(ctx.name),
+            line=ctx.call.line,
+            column=ctx.call.column,
+        )
+
+        ctx.api.add_symbol_table_node(
+            ctx.name, SymbolTableNode(GDEF, type_alias, plugin_generated=False)
+        )
+
+
 class StrawberryPlugin(Plugin):
+    def get_dynamic_class_hook(
+        self, fullname: str
+    ) -> Optional[Callable[[DynamicClassDefContext], None]]:
+        # TODO: investigate why we need this instead of `strawberry.union.union` on CI
+        # we have the same issue in the other hooks
+        if "strawberry.union" in fullname:
+            return union_hook
+
+        return None
+
     def get_type_analyze_hook(self, fullname: str):
         if fullname == "strawberry.lazy_type.LazyType":
             return lazy_type_analyze_callback


Test patch:
diff --git a/tests/mypy/test_plugin.yml b/tests/mypy/test_decorators.yml
similarity index 50%
rename from tests/mypy/test_plugin.yml
rename to tests/mypy/test_decorators.yml
--- a/tests/mypy/test_plugin.yml
+++ b/tests/mypy/test_decorators.yml
@@ -36,36 +36,3 @@
     NameInterface(n="Patrick")
   out: |
     main:8: error: Unexpected keyword argument "n" for "NameInterface"
-
-- case: test_field
-  main: |
-    import strawberry
-
-    @strawberry.type
-    class User:
-        name: str = strawberry.field(description='Example')
-
-    User(name="Patrick")
-    User(n="Patrick")
-  out: |
-    main:8: error: Unexpected keyword argument "n" for "User"
-
-- case: test_private_field
-  main: |
-    import strawberry
-
-    @strawberry.type
-    class User:
-        age: strawberry.Private[int]
-
-        @strawberry.field
-        def age_in_months(self) -> int:
-            return self.age * 12
-
-        @strawberry.field
-        def wrong_type(self) -> int:
-            reveal_type(self.age)
-            return self.age.trim()
-  out: |
-    main:13: note: Revealed type is 'builtins.int'
-    main:14: error: "int" has no attribute "trim"
diff --git a/tests/mypy/test_fields.yml b/tests/mypy/test_fields.yml
new file mode 100644
--- /dev/null
+++ b/tests/mypy/test_fields.yml
@@ -0,0 +1,32 @@
+- case: test_field
+  main: |
+    import strawberry
+
+    @strawberry.type
+    class User:
+        name: str = strawberry.field(description='Example')
+
+    User(name="Patrick")
+    User(n="Patrick")
+  out: |
+    main:8: error: Unexpected keyword argument "n" for "User"
+
+- case: test_private_field
+  main: |
+    import strawberry
+
+    @strawberry.type
+    class User:
+        age: strawberry.Private[int]
+
+        @strawberry.field
+        def age_in_months(self) -> int:
+            return self.age * 12
+
+        @strawberry.field
+        def wrong_type(self) -> int:
+            reveal_type(self.age)
+            return self.age.trim()
+  out: |
+    main:13: note: Revealed type is 'builtins.int'
+    main:14: error: "int" has no attribute "trim"
diff --git a/tests/mypy/test_union.yml b/tests/mypy/test_union.yml
new file mode 100644
--- /dev/null
+++ b/tests/mypy/test_union.yml
@@ -0,0 +1,51 @@
+- case: test_union
+  main: |
+    import typing
+
+    import strawberry
+
+    @strawberry.type
+    class User:
+        name: str
+
+    @strawberry.type
+    class Error:
+        message: str
+
+    Response = strawberry.union("Response", (User, Error))
+
+    a: Response
+    reveal_type(Response)
+    reveal_type(a)
+
+    a = User('abc')
+    reveal_type(a)
+  out: |
+    main:16: note: Revealed type is 'builtins.object'
+    main:17: note: Revealed type is 'Union[main.User, main.Error]'
+    main:20: note: Revealed type is 'main.User'
+
+- case: test_union_generics
+  main: |
+    from typing import Generic, TypeVar
+
+    import strawberry
+
+    T = TypeVar("T")
+
+    @strawberry.type
+    class Error:
+        message: str
+
+    @strawberry.type
+    class Edge(Generic[T]):
+        node: T
+
+    Result = strawberry.union("Result", (Error, Edge[str]))
+    reveal_type(Result)
+
+    a: Result
+    reveal_type(a)
+  out: |
+    main:16: note: Revealed type is 'builtins.object'
+    main:19: note: Revealed type is 'Union[main.Error, main.Edge[builtins.str]]'