From 0ed2ca59ad41f85f12fcd1fcd0564dc019d87a63 Mon Sep 17 00:00:00 2001
From: Adrien Vannson <adrien.vannson@protonmail.com>
Date: Fri, 3 Jan 2025 18:48:25 +0100
Subject: [PATCH] Add oneof documentation to generated code

---
 src/betterproto2_compiler/plugin/models.py    | 22 +++++++++++++++++++
 src/betterproto2_compiler/plugin/parser.py    | 19 +++++++++++++++-
 .../templates/template.py.j2                  |  9 +++++++-
 .../inputs/documentation/documentation.proto  | 14 ++++++++++++
 4 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/src/betterproto2_compiler/plugin/models.py b/src/betterproto2_compiler/plugin/models.py
index d71f75bd..0f95ae08 100644
--- a/src/betterproto2_compiler/plugin/models.py
+++ b/src/betterproto2_compiler/plugin/models.py
@@ -56,6 +56,7 @@
     FieldDescriptorProtoType as FieldType,
     FileDescriptorProto,
     MethodDescriptorProto,
+    OneofDescriptorProto,
 )
 from betterproto2_compiler.lib.google.protobuf.compiler import CodeGeneratorRequest
 
@@ -287,6 +288,7 @@ class MessageCompiler(ProtoContentBase):
     proto_obj: DescriptorProto = PLACEHOLDER
     path: list[int] = PLACEHOLDER
     fields: list[Union["FieldCompiler", "MessageCompiler"]] = field(default_factory=list)
+    oneofs: list["OneofCompiler"] = field(default_factory=list)
     builtins_types: set[str] = field(default_factory=set)
 
     def __post_init__(self) -> None:
@@ -566,6 +568,26 @@ def repeated(self) -> bool:
         return False  # maps cannot be repeated
 
 
+@dataclass
+class OneofCompiler(ProtoContentBase):
+    source_file: FileDescriptorProto
+    typing_compiler: TypingCompiler
+    path: list[int] = PLACEHOLDER
+
+    parent: MessageCompiler = PLACEHOLDER
+    proto_obj: OneofDescriptorProto = PLACEHOLDER
+
+    def __post_init__(self) -> None:
+        # Add oneof to message
+        if isinstance(self.parent, MessageCompiler):  # TODO why?
+            self.parent.oneofs.append(self)
+        super().__post_init__()
+
+    @property
+    def name(self) -> str:
+        return self.proto_obj.name
+
+
 @dataclass
 class EnumDefinitionCompiler(MessageCompiler):
     """Representation of a proto Enum definition."""
diff --git a/src/betterproto2_compiler/plugin/parser.py b/src/betterproto2_compiler/plugin/parser.py
index a2d9cbb3..740374f6 100644
--- a/src/betterproto2_compiler/plugin/parser.py
+++ b/src/betterproto2_compiler/plugin/parser.py
@@ -22,6 +22,7 @@
     FieldCompiler,
     MapEntryCompiler,
     MessageCompiler,
+    OneofCompiler,
     OneOfFieldCompiler,
     OutputTemplate,
     PluginRequestCompiler,
@@ -217,7 +218,13 @@ def read_protobuf_type(
                     typing_compiler=output_package.typing_compiler,
                 )
             elif is_oneof(field):
-                _make_one_of_field_compiler(output_package, source_file, message_data, field, path + [2, index])
+                OneOfFieldCompiler(
+                    source_file=source_file,
+                    parent=message_data,
+                    proto_obj=field,
+                    path=path + [2, index],
+                    typing_compiler=output_package.typing_compiler,
+                )
             else:
                 FieldCompiler(
                     source_file=source_file,
@@ -226,6 +233,16 @@ def read_protobuf_type(
                     path=path + [2, index],
                     typing_compiler=output_package.typing_compiler,
                 )
+
+        for index, oneof in enumerate(item.oneof_decl):
+            OneofCompiler(
+                source_file=source_file,
+                typing_compiler=output_package.typing_compiler,
+                path=path + [8, index],
+                parent=message_data,
+                proto_obj=oneof,
+            )
+
     elif isinstance(item, EnumDescriptorProto):
         # Enum
         EnumDefinitionCompiler(
diff --git a/src/betterproto2_compiler/templates/template.py.j2 b/src/betterproto2_compiler/templates/template.py.j2
index 3db14c37..af5d4b0a 100644
--- a/src/betterproto2_compiler/templates/template.py.j2
+++ b/src/betterproto2_compiler/templates/template.py.j2
@@ -32,9 +32,16 @@ class {{ enum.py_name }}(betterproto2.Enum):
 @dataclass(eq=False, repr=False)
 {% endif %}
 class {{ message.py_name }}(betterproto2.Message):
-    {% if message.comment %}
+    {% if message.comment or message.oneofs %}
     """
     {{ message.comment | indent(4) }}
+
+    {% if message.oneofs %}
+    Oneofs:
+        {% for oneof in message.oneofs %}
+        - {{ oneof.name }}: {{ oneof.comment | indent(12) }}
+        {% endfor %}
+    {% endif %}
     """
     {% endif %}
 
diff --git a/tests/inputs/documentation/documentation.proto b/tests/inputs/documentation/documentation.proto
index 7fc6c83b..8bb10ea0 100644
--- a/tests/inputs/documentation/documentation.proto
+++ b/tests/inputs/documentation/documentation.proto
@@ -13,6 +13,20 @@ message Test { // Documentation of message 3
     // Documentation of field 2
     // other line 2
     uint32 x = 1; // Documentation of field 3
+
+    // Documentation of oneof 1
+    // other line 1
+
+    // Documentation of oneof 2
+    // other line 2
+    oneof oneof_example { // Documentation of oneof 3
+        // Documentation of oneof field 1
+        // other line 1
+
+        // Documentation of oneof field 2
+        // other line 2
+        int32 a = 2; // Documentation of oneof field 3
+    }
 }
 
 // Documentation of enum 1