Repository: mirumee/ariadne-codegen (python)

Issue:
Handle query and mutation types names 
Reported in #100 
Given schema with query type with a different name than `Query` (the same for mutation type and name `Mutation`), eg.
```gql
schema {
    query: query_type
}

type query_type {
    test: String!
}
```
and operation
```gql
query Abc {
    test
}
```
`ariadne-codegen` doesn't generate client but raises exception instead, eg. `ariadne_codegen.exceptions.ParsingError: Field test not found in type Query.`

Pull request description:
Fix support of custom operation types names.
fixes #101 

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f3c9c35..69a39c8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,7 @@
 
 - Added generation of GraphQL schema's Python representation.
 - Fixed annotations for lists.
+- Fixed support of custom operation types names.
 
 
 ## 0.4.0 (2023-03-20)
diff --git a/ariadne_codegen/client_generators/result_types.py b/ariadne_codegen/client_generators/result_types.py
index 5336e61..dd7d4bc 100644
--- a/ariadne_codegen/client_generators/result_types.py
+++ b/ariadne_codegen/client_generators/result_types.py
@@ -15,6 +15,7 @@ from graphql import (
     InlineFragmentNode,
     NameNode,
     OperationDefinitionNode,
+    OperationType,
     SelectionNode,
     SelectionSetNode,
     StringValueNode,
@@ -92,10 +93,21 @@ class ResultTypesGenerator:
 
         self._class_defs = self._parse_type_definition(
             class_name=str_to_pascal_case(self.operation_definition.name.value),
-            type_name=self.operation_definition.operation.value.capitalize(),
+            type_name=self._get_operation_type_name(
+                self.operation_definition.operation
+            ),
             selection_set=self.operation_definition.selection_set,
         )
 
+    def _get_operation_type_name(self, operation_type: OperationType) -> str:
+        if operation_type == OperationType.QUERY and self.schema.query_type:
+            return self.schema.query_type.name
+
+        if operation_type == OperationType.MUTATION and self.schema.mutation_type:
+            return self.schema.mutation_type.name
+
+        raise NotSupported(f"Not supported operation type: {operation_type}")
+
     def generate(self) -> ast.Module:
         if self._used_enums:
             self._imports.append(


Test patch:
diff --git a/tests/client_generators/result_types_generator/test_parsing_operations.py b/tests/client_generators/result_types_generator/test_parsing_operations.py
index 3efaafa..2f06111 100644
--- a/tests/client_generators/result_types_generator/test_parsing_operations.py
+++ b/tests/client_generators/result_types_generator/test_parsing_operations.py
@@ -2,12 +2,7 @@ import ast
 from typing import Tuple, cast
 
 import pytest
-from graphql import (
-    FragmentDefinitionNode,
-    OperationDefinitionNode,
-    build_ast_schema,
-    parse,
-)
+from graphql import FragmentDefinitionNode, OperationDefinitionNode, build_schema, parse
 
 from ariadne_codegen.client_generators.constants import (
     BASE_MODEL_CLASS_NAME,
@@ -162,7 +157,7 @@ def test_generate_returns_module_with_generated_result_types_from_query(
         OperationDefinitionNode, parse(query_str).definitions[0]
     )
     generator = ResultTypesGenerator(
-        schema=build_ast_schema(parse(SCHEMA_STR)),
+        schema=build_schema(SCHEMA_STR),
         operation_definition=operation_definition,
         enums_module_name="enums",
     )
@@ -217,7 +212,7 @@ def test_generate_returns_module_with_types_generated_from_mutation():
         ),
     ]
     generator = ResultTypesGenerator(
-        schema=build_ast_schema(parse(SCHEMA_STR)),
+        schema=build_schema(SCHEMA_STR),
         operation_definition=operation_definition,
         enums_module_name="enums",
     )
@@ -255,7 +250,7 @@ def test_generate_returns_module_with_class_per_every_field():
         "CustomQueryQuery3Field2",
     ]
     generator = ResultTypesGenerator(
-        schema=build_ast_schema(parse(SCHEMA_STR)),
+        schema=build_schema(SCHEMA_STR),
         operation_definition=operation_definition,
         enums_module_name="enums",
     )
@@ -314,7 +309,7 @@ def test_generate_returns_module_with_types_generated_from_query_that_uses_fragm
         decorator_list=[],
     )
     generator = ResultTypesGenerator(
-        schema=build_ast_schema(parse(SCHEMA_STR)),
+        schema=build_schema(SCHEMA_STR),
         operation_definition=operation_definition,
         enums_module_name="enums",
         fragments_definitions={"TestFragment": fragment_def},
@@ -397,7 +392,7 @@ def test_generate_returns_module_with_class_for_every_appearance_of_type():
         ),
     ]
     generator = ResultTypesGenerator(
-        schema=build_ast_schema(parse(SCHEMA_STR)),
+        schema=build_schema(SCHEMA_STR),
         operation_definition=cast(
             OperationDefinitionNode, parse(query_str).definitions[0]
         ),
@@ -408,3 +403,68 @@ def test_generate_returns_module_with_class_for_every_appearance_of_type():
 
     class_defs = filter_class_defs(module)
     assert compare_ast(class_defs, expected_class_defs)
+
+
+@pytest.mark.parametrize(
+    "operation_str, expected_class_def",
+    [
+        (
+            "query TestQuery { testQ }",
+            ast.ClassDef(
+                name="TestQuery",
+                bases=[ast.Name(id=BASE_MODEL_CLASS_NAME)],
+                decorator_list=[],
+                keywords=[],
+                body=[
+                    ast.AnnAssign(
+                        target=ast.Name(id="testQ"),
+                        annotation=ast.Name(id="str"),
+                        simple=1,
+                    )
+                ],
+            ),
+        ),
+        (
+            "mutation TestMutation { testM }",
+            ast.ClassDef(
+                name="TestMutation",
+                bases=[ast.Name(id=BASE_MODEL_CLASS_NAME)],
+                decorator_list=[],
+                keywords=[],
+                body=[
+                    ast.AnnAssign(
+                        target=ast.Name(id="testM"),
+                        annotation=ast.Name(id="str"),
+                        simple=1,
+                    )
+                ],
+            ),
+        ),
+    ],
+)
+def test_generate_returns_module_for_schema_with_custom_operation_type_name(
+    operation_str, expected_class_def
+):
+    schema_str = """
+    schema {
+        query: CustomQueryType
+        mutation: CustomMutationType
+    }
+    type CustomQueryType { testQ: String! }
+    type CustomMutationType { testM: String! }
+    """
+    generator = ResultTypesGenerator(
+        schema=build_schema(schema_str),
+        operation_definition=cast(
+            OperationDefinitionNode, parse(operation_str).definitions[0]
+        ),
+        enums_module_name="enums",
+        convert_to_snake_case=False,
+    )
+
+    module = generator.generate()
+
+    assert isinstance(module, ast.Module)
+    class_defs = filter_class_defs(module)
+    assert len(class_defs) == 1
+    assert compare_ast(class_defs[0], expected_class_def)