Repository: strawberry-graphql/strawberry
Pull request #110 (2019-07-18 16:23:22)

Fix Union usage as return type for a mutation

codecov[bot]: # [Codecov](https://codecov.io/gh/strawberry-graphql/strawberry/pull/110?src=pr&el=h1) Report
> Merging [#110](https://codecov.io/gh/strawberry-graphql/strawberry/pull/110?src=pr&el=desc) into [master](https://codecov.io/gh/strawberry-graphql/strawberry/commit/1c5ba1e1d376d73f58c1be5662232641693f76a6?src=pr&el=desc) will **increase** coverage by `0.07%`.
> The diff coverage is `100%`.


```diff
@@            Coverage Diff             @@
##           master     #110      +/-   ##
==========================================
+ Coverage   90.02%   90.09%   +0.07%     
==========================================
  Files          20       20              
  Lines         401      404       +3     
==========================================
+ Hits          361      364       +3     
  Misses         40       40
```

patrick91: @marcoacierno let's add the release file and then we can merge this :)

Related to #4 

Patch:
diff --git a/RELEASE.md b/RELEASE.md
new file mode 100644
index 00000000..56cc37a6
--- /dev/null
+++ b/RELEASE.md
@@ -0,0 +1,37 @@
+Release type: patch
+
+Allow the usage of Union types in the mutations
+
+```python
+@strawberry.type
+class A:
+    x: int
+
+@strawberry.type
+class B:
+    y: int
+
+@strawberry.type
+class Mutation:
+    @strawberry.mutation
+    def hello(self, info) -> Union[A, B]:
+        return B(y=5)
+
+schema = strawberry.Schema(query=A, mutation=Mutation)
+
+query = """
+mutation {
+    hello {
+        __typename
+
+        ... on A {
+            x
+        }
+
+        ... on B {
+            y
+        }
+    }
+}
+"""
+```
diff --git a/strawberry/exceptions.py b/strawberry/exceptions.py
index 2f30f20d..e6b92bb5 100644
--- a/strawberry/exceptions.py
+++ b/strawberry/exceptions.py
@@ -2,6 +2,8 @@
 
 from typing import List, Set
 
+from graphql import GraphQLObjectType
+
 
 class NotAnEnum(Exception):
     def __init__(self):
@@ -40,3 +42,31 @@ class MissingArgumentsAnnotationsError(Exception):
         )
 
         super().__init__(message)
+
+
+class WrongReturnTypeForUnion(Exception):
+    """The Union type cannot be resolved because it's not a field"""
+
+    def __init__(self, field_name: str, result_type: str):
+        message = (
+            f'The type "{result_type}" cannot be resolved for the field "{field_name}" '
+            ", are you using a strawberry.field?"
+        )
+
+        super().__init__(message)
+
+
+class UnallowedReturnTypeForUnion(Exception):
+    """The return type is not in the list of Union types"""
+
+    def __init__(
+        self, field_name: str, result_type: str, allowed_types: Set[GraphQLObjectType]
+    ):
+        formatted_allowed_types = [type_.name for type_ in allowed_types]
+
+        message = (
+            f'The type "{result_type}" of the field "{field_name}" '
+            f'is not in the list of the types of the union: "{formatted_allowed_types}"'
+        )
+
+        super().__init__(message)
diff --git a/strawberry/type_converter.py b/strawberry/type_converter.py
index 2a378dce..0340e429 100644
--- a/strawberry/type_converter.py
+++ b/strawberry/type_converter.py
@@ -11,6 +11,7 @@ from graphql import (
     GraphQLUnionType,
 )
 
+from .exceptions import UnallowedReturnTypeForUnion, WrongReturnTypeForUnion
 from .scalars import ID
 from .utils.typing import is_union
 
@@ -65,6 +66,17 @@ def get_graphql_type_for_annotation(
             else:
                 is_field_optional = None.__class__ in types
 
+                def _resolve_type(self, value, _type):
+                    if not hasattr(self, "field"):
+                        raise WrongReturnTypeForUnion(value.field_name, str(type(self)))
+
+                    if self.field not in _type.types:
+                        raise UnallowedReturnTypeForUnion(
+                            value.field_name, str(type(self)), _type.types
+                        )
+
+                    return self.field
+
                 # TODO: union types don't work with scalar types
                 # so we want to return a nice error
                 # also we want to make sure we have been passed
@@ -72,6 +84,7 @@ def get_graphql_type_for_annotation(
                 graphql_type = GraphQLUnionType(
                     field_name, [type.field for type in types]
                 )
+                graphql_type.resolve_type = _resolve_type
         else:
             graphql_type = REGISTRY.get(annotation)
 


Test patch:
diff --git a/tests/test_union.py b/tests/test_union.py
new file mode 100644
index 00000000..2197d696
--- /dev/null
+++ b/tests/test_union.py
@@ -0,0 +1,155 @@
+from typing import Union
+
+import strawberry
+from graphql import graphql_sync
+
+
+def test_union_as_field():
+    @strawberry.type
+    class A:
+        a: int
+
+    @strawberry.type
+    class B:
+        b: int
+
+    @strawberry.type
+    class Query:
+        ab: Union[A, B] = A(a=5)
+
+    schema = strawberry.Schema(query=Query)
+    query = """{
+        ab {
+            __typename,
+
+            ... on A {
+                a
+            }
+        }
+    }"""
+
+    result = graphql_sync(schema, query)
+
+    assert not result.errors
+    assert result.data["ab"] == {"__typename": "A", "a": 5}
+
+
+def test_cannot_use_non_strawberry_fields_for_the_union():
+    @strawberry.type
+    class A:
+        a: int
+
+    @strawberry.type
+    class B:
+        b: int
+
+    @strawberry.type
+    class Query:
+        ab: Union[A, B] = "ciao"
+
+    schema = strawberry.Schema(query=Query)
+    query = """{
+        ab {
+            __typename,
+
+            ... on A {
+                a
+            }
+        }
+    }"""
+
+    result = graphql_sync(schema, query)
+
+    assert (
+        result.errors[0].message
+        == 'The type "<class \'str\'>" cannot be resolved for the field "ab" '
+        ", are you using a strawberry.field?"
+    )
+
+
+def test_union_as_mutation_return():
+    @strawberry.type
+    class A:
+        x: int
+
+    @strawberry.type
+    class B:
+        y: int
+
+    @strawberry.type
+    class Mutation:
+        @strawberry.mutation
+        def hello(self, info) -> Union[A, B]:
+            # TODO: mypy is unable to understand that B is a dataclass
+            return B(y=5)  # type:ignore
+
+    schema = strawberry.Schema(query=A, mutation=Mutation)
+
+    query = """
+    mutation {
+        hello {
+            __typename
+
+            ... on A {
+                x
+            }
+
+            ... on B {
+                y
+            }
+        }
+    }
+    """
+
+    result = graphql_sync(schema, query)
+
+    assert not result.errors
+    assert result.data["hello"] == {"__typename": "B", "y": 5}
+
+
+def test_types_not_included_in_the_union_are_rejected():
+    @strawberry.type
+    class Outside:
+        c: int
+
+    @strawberry.type
+    class A:
+        a: int
+
+    @strawberry.type
+    class B:
+        b: int
+
+    @strawberry.type
+    class Mutation:
+        @strawberry.mutation
+        def hello(self, info) -> Union[A, B]:
+            # TODO: mypy is unable to understand that Outside is a dataclass
+            return Outside(c=5)  # type:ignore
+
+    schema = strawberry.Schema(query=A, mutation=Mutation)
+
+    query = """
+    mutation {
+        hello {
+            __typename
+
+            ... on A {
+                a
+            }
+
+            ... on B {
+                b
+            }
+        }
+    }
+    """
+
+    result = graphql_sync(schema, query)
+
+    assert (
+        result.errors[0].message == "The type "
+        "\"<class 'tests.test_union.test_types_not_included_in_the_union_are_rejected.<locals>.Outside'>\""  # noqa
+        ' of the field "hello" '
+        "is not in the list of the types of the union: \"['A', 'B']\""
+    )