From c621ef8a8d4661e3a519e69a734eecd05127bed4 Mon Sep 17 00:00:00 2001
From: Adrien <adrien.vannson@protonmail.com>
Date: Tue, 22 Oct 2024 20:52:11 +0200
Subject: [PATCH] Fix documentation (#630)

* Fix missing documentation

* Add test

* Add test

* Format

* Reformat
---
 src/betterproto/plugin/models.py              |  2 ++
 src/betterproto/plugin/parser.py              | 17 ++++++++---
 .../inputs/documentation/documentation.proto  | 20 +++++++++++++
 tests/test_documentation.py                   | 29 +++++++++++++++++++
 4 files changed, 64 insertions(+), 4 deletions(-)
 create mode 100644 tests/inputs/documentation/documentation.proto
 create mode 100644 tests/test_documentation.py

diff --git a/src/betterproto/plugin/models.py b/src/betterproto/plugin/models.py
index 6b1e006..51735f4 100644
--- a/src/betterproto/plugin/models.py
+++ b/src/betterproto/plugin/models.py
@@ -661,6 +661,7 @@ class EnumDefinitionCompiler(MessageCompiler):
 
 @dataclass
 class ServiceCompiler(ProtoContentBase):
+    source_file: FileDescriptorProto
     parent: OutputTemplate = PLACEHOLDER
     proto_obj: DescriptorProto = PLACEHOLDER
     path: List[int] = PLACEHOLDER
@@ -682,6 +683,7 @@ class ServiceCompiler(ProtoContentBase):
 
 @dataclass
 class ServiceMethodCompiler(ProtoContentBase):
+    source_file: FileDescriptorProto
     parent: ServiceCompiler
     proto_obj: MethodDescriptorProto
     path: List[int] = PLACEHOLDER
diff --git a/src/betterproto/plugin/parser.py b/src/betterproto/plugin/parser.py
index 2e0d861..5f7b72c 100644
--- a/src/betterproto/plugin/parser.py
+++ b/src/betterproto/plugin/parser.py
@@ -143,7 +143,7 @@ def generate_code(request: CodeGeneratorRequest) -> CodeGeneratorResponse:
     for output_package_name, output_package in request_data.output_packages.items():
         for proto_input_file in output_package.input_files:
             for index, service in enumerate(proto_input_file.service):
-                read_protobuf_service(service, index, output_package)
+                read_protobuf_service(proto_input_file, service, index, output_package)
 
     # Generate output files
     output_paths: Set[pathlib.Path] = set()
@@ -249,12 +249,21 @@ def read_protobuf_type(
 
 
 def read_protobuf_service(
-    service: ServiceDescriptorProto, index: int, output_package: OutputTemplate
+    source_file: FileDescriptorProto,
+    service: ServiceDescriptorProto,
+    index: int,
+    output_package: OutputTemplate,
 ) -> None:
     service_data = ServiceCompiler(
-        parent=output_package, proto_obj=service, path=[6, index]
+        source_file=source_file,
+        parent=output_package,
+        proto_obj=service,
+        path=[6, index],
     )
     for j, method in enumerate(service.method):
         ServiceMethodCompiler(
-            parent=service_data, proto_obj=method, path=[6, index, 2, j]
+            source_file=source_file,
+            parent=service_data,
+            proto_obj=method,
+            path=[6, index, 2, j],
         )
diff --git a/tests/inputs/documentation/documentation.proto b/tests/inputs/documentation/documentation.proto
new file mode 100644
index 0000000..122785b
--- /dev/null
+++ b/tests/inputs/documentation/documentation.proto
@@ -0,0 +1,20 @@
+syntax = "proto3";
+package documentation;
+
+// Documentation of message
+message Test {
+    // Documentation of field
+    uint32 x = 1;
+}
+
+// Documentation of enum
+enum Enum {
+    // Documentation of variant
+    Enum_Variant = 0;
+}
+
+// Documentation of service
+service Service {
+    // Documentation of method
+    rpc get(Test) returns (Test);
+}
diff --git a/tests/test_documentation.py b/tests/test_documentation.py
new file mode 100644
index 0000000..78297bc
--- /dev/null
+++ b/tests/test_documentation.py
@@ -0,0 +1,29 @@
+import ast
+import inspect
+
+
+def test_documentation():
+    from .output_betterproto.documentation import (
+        Enum,
+        ServiceBase,
+        ServiceStub,
+        Test,
+    )
+
+    assert Test.__doc__ == "Documentation of message"
+
+    source = inspect.getsource(Test)
+    tree = ast.parse(source)
+    assert tree.body[0].body[2].value.value == "Documentation of field"
+
+    assert Enum.__doc__ == "Documentation of enum"
+
+    source = inspect.getsource(Enum)
+    tree = ast.parse(source)
+    assert tree.body[0].body[2].value.value == "Documentation of variant"
+
+    assert ServiceBase.__doc__ == "Documentation of service"
+    assert ServiceBase.get.__doc__ == "Documentation of method"
+
+    assert ServiceStub.__doc__ == "Documentation of service"
+    assert ServiceStub.get.__doc__ == "Documentation of method"