Take oneOf directive into account in codegen module

RELEASE.md

@@ -0,0 +1,48 @@
+Release type: minor
+
+The Query Codegen system now supports the `@oneOf` directive.
+
+When writing plugins, you can now access `GraphQLObjectType.is_one_of` to determine if the object being worked with has a `@oneOf` directive.
+
+The default plugins have been updated to take advantage of the new attribute.
+
+For example, given this schema:
+
+```python
+@strawberry.input(one_of=True)
+class OneOfInput:
+    a: Optional[str] = strawberry.UNSET
+    b: Optional[str] = strawberry.UNSET
+
+
+@strawberry.type
+class Query:
+    @strawberry.field
+    def one_of(self, value: OneOfInput) -> str: ...
+
+
+schema = strawberry.Schema(Query)
+```
+
+And this query:
+
+```graphql
+query OneOfTest($value: OneOfInput!) {
+  oneOf(value: $value)
+}
+```
+
+The query codegen can now generate this Typescript file:
+
+```typescript
+type OneOfTestResult = {
+    one_of: string
+}
+
+type OneOfInput = { a: string, b?: never }
+    | { a?: never, b: string }
+
+type OneOfTestVariables = {
+    value: OneOfInput
+}
+```

strawberry/codegen/plugins/python.py

@@ -109,16 +109,25 @@ def _get_type_name(self, type_: GraphQLType) -> str:
 
         return type_.name
 
-    def _print_field(self, field: GraphQLField) -> str:
+    def _print_field(self, field: GraphQLField, as_oneof_member: bool = False) -> str:
+        # `as_oneof_member` makes the field non optional
+        # We're doing this because we're expressing oneOf via union instead
+
         name = field.name
 
         if field.alias:
             name = f"# alias for {field.name}\n{field.alias}"
 
         default_value = ""
-        if field.default_value is not None:
+        if field.default_value is not None and not as_oneof_member:
             default_value = f" = {self._print_argument_value(field.default_value)}"
-        return f"{name}: {self._get_type_name(field.type)}{default_value}"
+
+        if as_oneof_member and isinstance(field.type, GraphQLOptional):
+            type_ = field.type.of_type
+        else:
+            type_ = field.type
+
+        return f"{name}: {self._get_type_name(type_)}{default_value}"
 
     def _print_argument_value(self, argval: GraphQLArgumentValue) -> str:
         if hasattr(argval, "values"):
@@ -174,6 +183,37 @@ def _print_object_type(self, type_: GraphQLObjectType) -> str:
 
         return "\n".join(lines)
 
+    def _get_oneof_class_name(
+        self, parent_type: GraphQLObjectType, member_field: GraphQLField
+    ) -> str:
+        # Name the classes using the parent name and field name
+        # Example.option => ExampleOption
+        return f"{parent_type.name}{member_field.name.title()}"
+
+    def _print_oneof_object_type(self, type_: GraphQLObjectType) -> str:
+        self.imports["typing"].add("Union")
+
+        fields = [field for field in type_.fields if field.name != "__typename"]
+
+        indent = 4 * " "
+
+        lines = []
+        for field in fields:
+            # Add a one-field class for each oneOf member
+            lines.append(f"class {self._get_oneof_class_name(type_, field)}:")
+            lines.append(
+                textwrap.indent(self._print_field(field, as_oneof_member=True), indent)
+            )
+            lines.append("")
+
+        # Create a union of the classes we just created
+        type_list = ", ".join(
+            [self._get_oneof_class_name(type_, field) for field in fields]
+        )
+        lines.append(f"{type_.name} = Union[{type_list}]")
+
+        return "\n".join(lines)
+
     def _print_enum_type(self, type_: GraphQLEnum) -> str:
         values = "\n".join(self._print_enum_value(value) for value in type_.values)
 
@@ -202,7 +242,10 @@ def _print_type(self, type_: GraphQLType) -> str:
             return self._print_union_type(type_)
 
         if isinstance(type_, GraphQLObjectType):
-            return self._print_object_type(type_)
+            if type_.is_one_of:
+                return self._print_oneof_object_type(type_)
+            else:
+                return self._print_object_type(type_)
 
         if isinstance(type_, GraphQLEnum):
             return self._print_enum_type(type_)

strawberry/codegen/plugins/typescript.py

@@ -77,6 +77,18 @@ def _print_field(self, field: GraphQLField) -> str:
 
         return f"{name}: {self._get_type_name(field.type)}"
 
+    def _print_oneof_field(self, field: GraphQLField) -> str:
+        name = field.name
+
+        if isinstance(field.type, GraphQLOptional):
+            # Use the non-null version of the type because we're using unions instead
+            output_type = field.type.of_type
+        else:
+            # Shouldn't run, oneOf types are always nullable
+            # Keeping it here just in case
+            output_type = field.type  # pragma: no cover
+        return f"{name}: {self._get_type_name(output_type)}"
+
     def _print_enum_value(self, value: str) -> str:
         return f'{value} = "{value}",'
 
@@ -87,6 +99,27 @@ def _print_object_type(self, type_: GraphQLObjectType) -> str:
             [f"type {type_.name} = {{", textwrap.indent(fields, " " * 4), "}"],
         )
 
+    def _print_oneof_object_type(self, type_: GraphQLObjectType) -> str:
+        # We'll gather a list of objects for each oneOf field
+        options: list[str] = []
+        for option in type_.fields:
+            # We'll give each option all fields from the parent type
+            option_fields: list[str] = []
+            for field in type_.fields:
+                if field == option:
+                    # Each option gets one field with its type...
+                    field_row = self._print_oneof_field(field)
+                else:
+                    # ... and the rest set to `never` to prevent multiple from being set
+                    field_row = f"{field.name}?: never"
+                option_fields.append(field_row)
+            options.append("{ " + ", ".join(option_fields) + " }")
+
+        # Union all the options together
+        all_options = "\n    | ".join(options)
+
+        return f"type {type_.name} = {all_options}"
+
     def _print_enum_type(self, type_: GraphQLEnum) -> str:
         values = "\n".join(self._print_enum_value(value) for value in type_.values)
 
@@ -112,7 +145,10 @@ def _print_type(self, type_: GraphQLType) -> str:
             return self._print_union_type(type_)
 
         if isinstance(type_, GraphQLObjectType):
-            return self._print_object_type(type_)
+            if type_.is_one_of:
+                return self._print_oneof_object_type(type_)
+            else:
+                return self._print_object_type(type_)
 
         if isinstance(type_, GraphQLEnum):
             return self._print_enum_type(type_)

strawberry/codegen/query_codegen.py

@@ -582,6 +582,7 @@ def _collect_type_from_strawberry_type(
             type_ = GraphQLObjectType(
                 strawberry_type.name,
                 [],
+                is_one_of=strawberry_type.is_one_of,
             )
 
             for field in strawberry_type.fields:

strawberry/codegen/types.py

@@ -44,6 +44,7 @@ class GraphQLObjectType:
     name: str
     fields: List[GraphQLField] = field(default_factory=list)
     graphql_typename: Optional[str] = None
+    is_one_of: bool = False
 
 
 # Subtype of GraphQLObjectType.

tests/codegen/conftest.py

@@ -83,6 +83,12 @@ class ExampleInput:
     optional_people: Optional[List[PersonInput]]
 
 
+@strawberry.input(one_of=True)
+class OneOfInput:
+    a: Optional[str] = strawberry.UNSET
+    b: Optional[str] = strawberry.UNSET
+
+
 @strawberry.type
 class Query:
     id: strawberry.ID
@@ -127,6 +133,14 @@ def list_life() -> LifeContainer[Person, Animal]:
         dinosaur = Animal(name="rex", age=66_000_000)
         return LifeContainer([person], [dinosaur])
 
+    @strawberry.field
+    def one_of(self, value: OneOfInput) -> str: ...  # pragma: no cover
+
+    @strawberry.field
+    def one_of_typename(
+        self, value: OneOfInput
+    ) -> PersonOrAnimal: ...  # pragma: no cover
+
 
 @strawberry.input
 class BlogPostInput:

tests/codegen/queries/oneof.graphql

@@ -0,0 +1,3 @@
+query OneOfTest($value: OneOfInput!) {
+  oneOf(value: $value)
+}

tests/codegen/queries/oneof_typename.graphql

@@ -0,0 +1,8 @@
+query OneOfTypenameTest($value: OneOfInput!) {
+  alias: oneOfTypename(value: $value) {
+    ... on Person {
+      name
+      age
+    }
+  }
+}

tests/codegen/snapshots/python/oneof.py

@@ -0,0 +1,15 @@
+from typing import Union
+
+class OneOfTestResult:
+    one_of: str
+
+class OneOfInputA:
+    a: str
+
+class OneOfInputB:
+    b: str
+
+OneOfInput = Union[OneOfInputA, OneOfInputB]
+
+class OneOfTestVariables:
+    value: OneOfInput

tests/codegen/snapshots/python/oneof_typename.py

@@ -0,0 +1,21 @@
+from typing import Union
+
+class OneOfTypenameTestResultOneOfTypenamePerson:
+    # typename: Person
+    name: str
+    age: int
+
+class OneOfTypenameTestResult:
+    # alias for one_of_typename
+    alias: OneOfTypenameTestResultOneOfTypenamePerson
+
+class OneOfInputA:
+    a: str
+
+class OneOfInputB:
+    b: str
+
+OneOfInput = Union[OneOfInputA, OneOfInputB]
+
+class OneOfTypenameTestVariables:
+    value: OneOfInput

tests/codegen/snapshots/typescript/oneof.ts

@@ -0,0 +1,10 @@
+type OneOfTestResult = {
+    one_of: string
+}
+
+type OneOfInput = { a: string, b?: never }
+    | { a?: never, b: string }
+
+type OneOfTestVariables = {
+    value: OneOfInput
+}

tests/codegen/snapshots/typescript/oneof_typename.ts

@@ -0,0 +1,16 @@
+type OneOfTypenameTestResultOneOfTypenamePerson = {
+    name: string
+    age: number
+}
+
+type OneOfTypenameTestResult = {
+    // alias for one_of_typename
+    alias: OneOfTypenameTestResultOneOfTypenamePerson
+}
+
+type OneOfInput = { a: string, b?: never }
+    | { a?: never, b: string }
+
+type OneOfTypenameTestVariables = {
+    value: OneOfInput
+}