Take oneOf directive into account in codegen module

strawberry/codegen/plugins/python.py

@@ -109,16 +109,22 @@ 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:
         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 +180,36 @@ 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:
+        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:
+            lines.append(f"class {self._get_oneof_class_name(type_, field)}:")
+            lines.append(
+                textwrap.indent(self._print_field(field, as_oneof_member=True), indent)
+            )
+
+        type_list = ",".join(
+            [self._get_oneof_class_name(type_, field) for field in fields]
+        )
+        lines.append(f"{type_.name} = Union[{type_list}]")
+
+        if type_.graphql_typename:
+            # Shouldn't run, inputs can't be fragments
+            lines.append(f"# typename: {type_.graphql_typename}")  # pragma: no cover
+
+        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 +238,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,20 @@ 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 field.alias:
+            # Shouldn't run, aliases can't exist on inputs
+            name = f"// alias for {field.name}\n{field.alias}"  # pragma: no cover
+
+        if isinstance(field.type, GraphQLOptional):
+            output_type = field.type.of_type
+        else:
+            # Shouldn't run, oneOf types are always nullable
+            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 +101,22 @@ 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:
+        options: list[str] = []
+        for option in type_.fields:
+            option_fields: list[str] = []
+            for field in type_.fields:
+                if field == option:
+                    field_row = self._print_oneof_field(field)
+                else:
+                    field_row = f"{field.name}: never"
+                option_fields.append(textwrap.indent(field_row, " " * 4))
+            options.append("{\n" + ",\n".join(option_fields) + "\n}")
+
+        all_options = " | ".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 +142,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,12 @@ 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: ...
+
+    @strawberry.field
+    def one_of_typename(self, value: OneOfInput) -> PersonOrAnimal: ...
+
 
 @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,13 @@
+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,19 @@
+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,15 @@
+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,21 @@
+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
+}