Add all kinds of comments (#632)

* Add all kinds of comments

* Format

* Fix get_comment

* Improve test

* Lint

src/betterproto/plugin/models.py

@@ -153,11 +153,33 @@ def get_comment(
 ) -> str:
     pad = " " * indent
     for sci_loc in proto_file.source_code_info.location:
-        if list(sci_loc.path) == path and sci_loc.leading_comments:
+        if list(sci_loc.path) == path:
-            lines = sci_loc.leading_comments.strip().split("\n")
+            all_comments = list(sci_loc.leading_detached_comments)
+            if sci_loc.leading_comments:
+                all_comments.append(sci_loc.leading_comments)
+            if sci_loc.trailing_comments:
+                all_comments.append(sci_loc.trailing_comments)
+
+            lines = []
+
+            for comment in all_comments:
+                lines += comment.split("\n")
+                lines.append("")
+
+            # Remove consecutive empty lines
+            lines = [
+                line for i, line in enumerate(lines) if line or (i == 0 or lines[i - 1])
+            ]
+
+            if lines and not lines[-1]:
+                lines.pop()  # Remove the last empty line
+
+            # It is common for one line comments to start with a space, for example: // comment
+            # We don't add this space to the generated file.
+            lines = [line[1:] if line and line[0] == " " else line for line in lines]
+
             # This is a field, message, enum, service, or method
             if len(lines) == 1 and len(lines[0]) < 79 - indent - 6:
-                lines[0] = lines[0].strip('"')
                 return f'{pad}"""{lines[0]}"""'
             else:
                 joined = f"\n{pad}".join(lines)

tests/inputs/documentation/documentation.proto

@@ -1,20 +1,44 @@
 syntax = "proto3";
 package documentation;
 
-// Documentation of message
+// Documentation of message 1
-message Test {
+// other line 1
-    // Documentation of field
+
-    uint32 x = 1;
+// Documentation of message 2
+// other line 2
+message Test { // Documentation of message 3
+    // Documentation of field 1
+    // other line 1
+
+    // Documentation of field 2
+    // other line 2
+    uint32 x = 1; // Documentation of field 3
 }
 
-// Documentation of enum
+// Documentation of enum 1
-enum Enum {
+// other line 1
-    // Documentation of variant
+
-    Enum_Variant = 0;
+// Documentation of enum 2
+// other line 2
+enum Enum { // Documentation of enum 3
+    // Documentation of variant 1
+    // other line 1
+
+    // Documentation of variant 2
+    // other line 2
+    Enum_Variant = 0; // Documentation of variant 3
 }
 
-// Documentation of service
+// Documentation of service 1
-service Service {
+// other line 1
-    // Documentation of method
+
-    rpc get(Test) returns (Test);
+// Documentation of service 2
+// other line 2
+service Service { // Documentation of service 3
+    // Documentation of method 1
+    // other line 1
+
+    // Documentation of method 2
+    // other line 2
+    rpc get(Test) returns (Test); // Documentation of method 3
 }

tests/test_documentation.py

@@ -2,7 +2,15 @@ import ast
 import inspect
 
 
-def test_documentation():
+def check(generated_doc: str, type: str) -> None:
+    assert f"Documentation of {type} 1" in generated_doc
+    assert "other line 1" in generated_doc
+    assert f"Documentation of {type} 2" in generated_doc
+    assert "other line 2" in generated_doc
+    assert f"Documentation of {type} 3" in generated_doc
+
+
+def test_documentation() -> None:
     from .output_betterproto.documentation import (
         Enum,
         ServiceBase,
@@ -10,20 +18,20 @@ def test_documentation():
         Test,
     )
 
-    assert Test.__doc__ == "Documentation of message"
+    check(Test.__doc__, "message")
 
     source = inspect.getsource(Test)
     tree = ast.parse(source)
-    assert tree.body[0].body[2].value.value == "Documentation of field"
+    check(tree.body[0].body[2].value.value, "field")
 
-    assert Enum.__doc__ == "Documentation of enum"
+    check(Enum.__doc__, "enum")
 
     source = inspect.getsource(Enum)
     tree = ast.parse(source)
-    assert tree.body[0].body[2].value.value == "Documentation of variant"
+    check(tree.body[0].body[2].value.value, "variant")
 
-    assert ServiceBase.__doc__ == "Documentation of service"
+    check(ServiceBase.__doc__, "service")
-    assert ServiceBase.get.__doc__ == "Documentation of method"
+    check(ServiceBase.get.__doc__, "method")
 
-    assert ServiceStub.__doc__ == "Documentation of service"
+    check(ServiceStub.__doc__, "service")
-    assert ServiceStub.get.__doc__ == "Documentation of method"
+    check(ServiceStub.get.__doc__, "method")