Add docstring to generated code based on schema description (#362)

* Add docstring to generated code based on schema description

* add method description

* add docstring for input field

* add field docstring to custom_fields class

* fix docstring of method

* add tests

* lint

---------

Co-authored-by: John Lyu <lvjunhong@citics.com>

Author: PaleNeutron

ariadne_codegen/client_generators/custom_fields.py

@@ -114,6 +114,7 @@ def _parse_object_type_definitions(
                 class_def = self._generate_class_def_body(
                     definition=graphql_type,
                     class_name=f"{graphql_type.name}{self._get_suffix(graphql_type)}",
+                    description=graphql_type.description,
                 )
                 if isinstance(graphql_type, GraphQLInterfaceType):
                     class_def.body.append(
@@ -129,17 +130,20 @@ def _generate_class_def_body(
         self,
         definition: Union[GraphQLObjectType, GraphQLInterfaceType],
         class_name: str,
+        description: Optional[str] = None,
     ) -> ast.ClassDef:
         """
         Generates the body of a class definition for a given GraphQL object
         or interface type.
         """
         base_names = [GRAPHQL_BASE_FIELD_CLASS]
         additional_fields_typing = set()
-        class_def = generate_class_def(name=class_name, base_names=base_names)
-        for lineno, (org_name, field) in enumerate(
-            self._get_combined_fields(definition).items(), start=1
-        ):
+        class_def = generate_class_def(
+            name=class_name, base_names=base_names, description=description
+        )
+        lineno = 0
+        for org_name, field in self._get_combined_fields(definition).items():
+            lineno += 1
             name = process_name(
                 org_name, convert_to_snake_case=self.convert_to_snake_case
             )
@@ -154,6 +158,11 @@ def _generate_class_def_body(
                     name, field_name, org_name, field, method_required, lineno
                 )
             )
+            # Add field docstring for class attributes (not methods)
+            if not getattr(field, "args") and field.description and not method_required:
+                lineno += 1
+                docstring = ast.Expr(value=ast.Constant(field.description))
+                class_def.body.append(docstring)
 
         class_def.body.append(
             self._generate_fields_method(
@@ -216,7 +225,11 @@ def _generate_class_field(
         """Handles the generation of field types."""
         if getattr(field, "args") or method_required:
             return self.generate_product_type_method(
-                name, field_name, org_name, getattr(field, "args")
+                name,
+                field_name,
+                org_name,
+                getattr(field, "args"),
+                description=getattr(field, "description"),
             )
         return generate_ann_assign(
             target=generate_name(name),
@@ -316,6 +329,7 @@ def generate_product_type_method(
         class_name: str,
         org_name: str,
         arguments: Optional[Dict[str, Any]] = None,
+        description: Optional[str] = None,
     ) -> ast.FunctionDef:
         """Generates a method for a product type."""
         arguments = arguments or {}
@@ -355,6 +369,7 @@ def generate_product_type_method(
             ),
             return_type=generate_name(f'"{class_name}"'),
             decorator_list=[generate_name("classmethod")],
+            description=description,
         )
 
     def _get_suffix(

ariadne_codegen/client_generators/custom_operation.py

@@ -84,6 +84,7 @@ def generate(self) -> ast.Module:
                 operation_name=name,
                 operation_args=field.args,
                 final_type=final_type,
+                description=field.description,
             )
             method_def.lineno = len(self._class_def.body) + 1
             self._class_def.body.append(method_def)
@@ -115,6 +116,7 @@ def _generate_method(
         operation_name: str,
         operation_args,
         final_type,
+        description: Optional[str] = None,
     ) -> ast.FunctionDef:
         """Generates a method definition for a given operation."""
         (
@@ -141,6 +143,7 @@ def _generate_method(
             name=str_to_snake_case(operation_name),
             arguments=method_arguments,
             return_type=generate_name(return_type_name),
+            description=description,
             body=[
                 *arguments_body,
                 generate_return(

ariadne_codegen/client_generators/input_types.py

@@ -157,10 +157,13 @@ def _parse_input_definition(
         self, definition: GraphQLInputObjectType
     ) -> ast.ClassDef:
         class_def = generate_class_def(
-            name=definition.name, base_names=[BASE_MODEL_CLASS_NAME]
+            name=definition.name,
+            base_names=[BASE_MODEL_CLASS_NAME],
+            description=definition.description,
         )
-
-        for lineno, (org_name, field) in enumerate(definition.fields.items(), start=1):
+        lineno = 0
+        for org_name, field in definition.fields.items():
+            lineno += 1
             name = process_name(
                 org_name,
                 convert_to_snake_case=self.convert_to_snake_case,
@@ -190,6 +193,10 @@ def _parse_input_definition(
                     field_implementation, input_field=field, field_name=org_name
                 )
             class_def.body.append(field_implementation)
+            if field.description:
+                lineno += 1
+                docstring = ast.Expr(value=ast.Constant(value=field.description))
+                class_def.body.append(docstring)
             self._save_dependencies(root_type=definition.name, field_type=field_type)
 
         if self.plugin_manager:

ariadne_codegen/codegen.py

@@ -111,16 +111,21 @@ def generate_class_def(
     name: str,
     base_names: Optional[List[str]] = None,
     body: Optional[List[ast.stmt]] = None,
+    description: str = "",
 ) -> ast.ClassDef:
     """Generate class definition."""
     bases = cast(
         List[ast.expr], [ast.Name(id=name) for name in base_names] if base_names else []
     )
+    body = body if body else []
+    if description:
+        docstring = ast.Expr(value=ast.Constant(value=description))
+        body.insert(0, docstring)
     params: Dict[str, Any] = {
         "name": name,
         "bases": bases,
         "keywords": [],
-        "body": body if body else [],
+        "body": body,
         "decorator_list": [],
     }
     if sys.version_info >= (3, 12):
@@ -354,10 +359,15 @@ def generate_method_definition(
     name: str,
     arguments: ast.arguments,
     return_type: Union[ast.Name, ast.Subscript],
+    description: str = "",
     body: Optional[List[ast.stmt]] = None,
     lineno: int = 1,
     decorator_list: Optional[List[ast.expr]] = None,
 ) -> ast.FunctionDef:
+    body = body if body else [ast.Pass()]
+    if description:
+        docstring = ast.Expr(value=ast.Constant(value=description))
+        body.insert(0, docstring)
     params: Dict[str, Any] = {
         "name": name,
         "args": arguments,

tests/codegen/test_generate_method_definition.py

@@ -0,0 +1,165 @@
+import ast
+from unittest import mock
+
+from ariadne_codegen.codegen import generate_arguments, generate_method_definition
+
+
+def test_generate_method_definition_with_minimal_parameters():
+    """Test generate_method_definition with only required parameters."""
+    name = "test_method"
+    arguments = generate_arguments()
+    return_type = ast.Name(id="str")
+
+    result = generate_method_definition(name, arguments, return_type)
+
+    assert isinstance(result, ast.FunctionDef)
+    assert result.name == name
+    assert result.args == arguments
+    assert result.returns == return_type
+    assert result.lineno == 1
+    assert len(result.body) == 1
+    assert isinstance(result.body[0], ast.Pass)
+    assert result.decorator_list == []
+
+
+def test_generate_method_definition_with_description():
+    """Test generate_method_definition with description adds docstring."""
+    name = "test_method"
+    arguments = generate_arguments()
+    return_type = ast.Name(id="str")
+    description = "This is a test method"
+
+    result = generate_method_definition(
+        name, arguments, return_type, description=description
+    )
+
+    assert len(result.body) == 2
+    assert isinstance(result.body[0], ast.Expr)
+    assert isinstance(result.body[0].value, ast.Constant)
+    assert result.body[0].value.value == description
+    assert isinstance(result.body[1], ast.Pass)
+
+
+def test_generate_method_definition_with_custom_body():
+    """Test generate_method_definition with custom body."""
+    name = "test_method"
+    arguments = generate_arguments()
+    return_type = ast.Name(id="str")
+    body = [ast.Return(value=ast.Constant(value="test"))]
+
+    result = generate_method_definition(name, arguments, return_type, body=body)
+
+    assert result.body == body
+    assert len(result.body) == 1
+    assert isinstance(result.body[0], ast.Return)
+
+
+def test_generate_method_definition_with_description_and_custom_body():
+    """Test generate_method_definition with both description and custom body."""
+    name = "test_method"
+    arguments = generate_arguments()
+    return_type = ast.Name(id="str")
+    description = "Test description"
+    body = [ast.Return(value=ast.Constant(value="test"))]
+
+    result = generate_method_definition(
+        name, arguments, return_type, description=description, body=body
+    )
+
+    assert len(result.body) == 2
+    assert isinstance(result.body[0], ast.Expr)
+    assert result.body[0].value.value == description
+    assert isinstance(result.body[1], ast.Return)
+
+
+def test_generate_method_definition_with_custom_lineno():
+    """Test generate_method_definition with custom line number."""
+    name = "test_method"
+    arguments = generate_arguments()
+    return_type = ast.Name(id="str")
+    lineno = 42
+
+    result = generate_method_definition(name, arguments, return_type, lineno=lineno)
+
+    assert result.lineno == lineno
+
+
+def test_generate_method_definition_with_decorators():
+    """Test generate_method_definition with decorator list."""
+    name = "test_method"
+    arguments = generate_arguments()
+    return_type = ast.Name(id="str")
+    decorators = [ast.Name(id="property"), ast.Name(id="staticmethod")]
+
+    result = generate_method_definition(
+        name, arguments, return_type, decorator_list=decorators
+    )
+
+    assert result.decorator_list == decorators
+
+
+def test_generate_method_definition_with_subscript_return_type():
+    """Test generate_method_definition with Subscript return type."""
+    name = "test_method"
+    arguments = generate_arguments()
+    return_type = ast.Subscript(value=ast.Name(id="List"), slice=ast.Name(id="str"))
+
+    result = generate_method_definition(name, arguments, return_type)
+
+    assert result.returns == return_type
+    assert isinstance(result.returns, ast.Subscript)
+
+
+@mock.patch("sys.version_info", (3, 12, 0))
+def test_generate_method_definition_python_312_or_later():
+    """Test generate_method_definition includes type_params for Python 3.12+."""
+    name = "test_method"
+    arguments = generate_arguments()
+    return_type = ast.Name(id="str")
+
+    result = generate_method_definition(name, arguments, return_type)
+
+    assert hasattr(result, "type_params")
+    assert result.type_params == []
+
+
+@mock.patch("sys.version_info", (3, 11, 0))
+def test_generate_method_definition_python_311_or_earlier():
+    """Test generate_method_definition doesn't include type_params for Python < 3.12."""
+    name = "test_method"
+    arguments = generate_arguments()
+    return_type = ast.Name(id="str")
+
+    result = generate_method_definition(name, arguments, return_type)
+
+    # In Python < 3.12, type_params shouldn't be set
+    # We check that the function still works correctly
+    assert isinstance(result, ast.FunctionDef)
+    assert result.name == name
+
+
+def test_generate_method_definition_all_parameters():
+    """Test generate_method_definition with all parameters specified."""
+    name = "complex_method"
+    arguments = generate_arguments()
+    return_type = ast.Subscript(value=ast.Name(id="Optional"), slice=ast.Name(id="int"))
+    description = "A complex method with all parameters"
+    body = [
+        ast.Assign(targets=[ast.Name(id="x")], value=ast.Constant(value=1)),
+        ast.Return(value=ast.Name(id="x")),
+    ]
+    lineno = 10
+    decorators = [ast.Name(id="classmethod")]
+
+    result = generate_method_definition(
+        name, arguments, return_type, description, body, lineno, decorators
+    )
+
+    assert result.name == name
+    assert result.args == arguments
+    assert result.returns == return_type
+    assert result.lineno == lineno
+    assert result.decorator_list == decorators
+    assert len(result.body) == 3  # docstring + 2 body statements
+    assert isinstance(result.body[0], ast.Expr)  # docstring
+    assert result.body[0].value.value == description

tests/codegen/test_generated_classes.py

@@ -24,3 +24,17 @@ def test_generate_class_def_returns_class_def_without_base():
     assert isinstance(result, ast.ClassDef)
     assert result.name == name
     assert not result.bases
+
+
+def test_generate_class_def_with_description_adds_docstring():
+    name = "Xyz"
+    description = "This is a test class. \nWith multiple lines."
+
+    result = generate_class_def(name, description=description)
+
+    assert isinstance(result, ast.ClassDef)
+    assert result.name == name
+    docstring = result.body[0]
+    assert isinstance(docstring, ast.Expr)
+    assert isinstance(docstring.value, ast.Constant)
+    assert docstring.value.value == description