From 93491a333c30bd9c7cc088cb64f95d6b3ab778fc Mon Sep 17 00:00:00 2001
From: Adrien Vannson <adrien.vannson@protonmail.com>
Date: Fri, 2 May 2025 20:31:24 +0200
Subject: [PATCH 1/2] Fix __all__ definition

---
 src/betterproto2_compiler/plugin/compiler.py       | 14 ++++++++++++--
 src/betterproto2_compiler/templates/header.py.j2   |  8 +++++++-
 .../templates/service_stub.py.j2                   |  2 +-
 src/betterproto2_compiler/templates/template.py.j2 |  6 +++---
 4 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/src/betterproto2_compiler/plugin/compiler.py b/src/betterproto2_compiler/plugin/compiler.py
index c35697e7..62f1658e 100644
--- a/src/betterproto2_compiler/plugin/compiler.py
+++ b/src/betterproto2_compiler/plugin/compiler.py
@@ -33,12 +33,22 @@ def outputfile_compiler(output_file: OutputTemplate) -> str:
         loader=jinja2.FileSystemLoader(templates_folder),
         undefined=jinja2.StrictUndefined,
     )
-    # Load the body first so we have a compleate list of imports needed.
+
+    # List of the symbols that should appear in the `__all__` variable of the file
+    all: list[str] = []
+
+    def add_to_all(name: str) -> str:
+        all.append(name)
+        return name
+
+    env.filters["add_to_all"] = add_to_all
+
     body_template = env.get_template("template.py.j2")
     header_template = env.get_template("header.py.j2")
 
+    # Load the body first do know the symbols defined in the file
     code = body_template.render(output_file=output_file)
-    code = header_template.render(output_file=output_file, version=version) + "\n" + code
+    code = header_template.render(output_file=output_file, version=version, all=all) + "\n" + code
 
     try:
         # Sort imports, delete unused ones
diff --git a/src/betterproto2_compiler/templates/header.py.j2 b/src/betterproto2_compiler/templates/header.py.j2
index d6bb310f..74738fcc 100644
--- a/src/betterproto2_compiler/templates/header.py.j2
+++ b/src/betterproto2_compiler/templates/header.py.j2
@@ -5,7 +5,7 @@
 # plugin: python-betterproto2
 # This file has been @generated
 
-__all__ = (
+{# __all__ = (
     {% for _, enum in output_file.enums|dictsort(by="key") %}
         "{{ enum.py_name }}",
     {%- endfor -%}
@@ -16,6 +16,12 @@ __all__ = (
         "{{ service.py_name }}Stub",
         "{{ service.py_name }}Base",
     {%- endfor -%}
+) #}
+
+__all__ = (
+    {%- for name in all -%}
+    "{{ name }}",
+    {%- endfor -%}
 )
 
 import re
diff --git a/src/betterproto2_compiler/templates/service_stub.py.j2 b/src/betterproto2_compiler/templates/service_stub.py.j2
index f0958055..54cec014 100644
--- a/src/betterproto2_compiler/templates/service_stub.py.j2
+++ b/src/betterproto2_compiler/templates/service_stub.py.j2
@@ -1,4 +1,4 @@
-class {% block class_name %}{% endblock %}({% block inherit_from %}{% endblock %}):
+class {% filter add_to_all %}{% block class_name %}{% endblock %}{% endfilter %}({% block inherit_from %}{% endblock %}):
     {% block service_docstring scoped %}
     {% if service.comment %}
     """
diff --git a/src/betterproto2_compiler/templates/template.py.j2 b/src/betterproto2_compiler/templates/template.py.j2
index 4f89d832..f51a3b39 100644
--- a/src/betterproto2_compiler/templates/template.py.j2
+++ b/src/betterproto2_compiler/templates/template.py.j2
@@ -1,5 +1,5 @@
 {% for _, enum in output_file.enums|dictsort(by="key") %}
-class {{ enum.py_name }}(betterproto2.Enum):
+class {{ enum.py_name | add_to_all }}(betterproto2.Enum):
     {% if enum.comment %}
     """
     {{ enum.comment | indent(4) }}
@@ -31,7 +31,7 @@ class {{ enum.py_name }}(betterproto2.Enum):
 {% else %}
 @dataclass(eq=False, repr=False)
 {% endif %}
-class {{ message.py_name }}(betterproto2.Message):
+class {{ message.py_name | add_to_all }}(betterproto2.Message):
     {% if message.comment or message.oneofs %}
     """
     {{ message.comment | indent(4) }}
@@ -104,7 +104,7 @@ default_message_pool.register_message("{{ output_file.package }}", "{{ message.p
 
 {% if output_file.settings.server_generation == "async" %}
 {% for _, service in output_file.services|dictsort(by="key") %}
-class {{ service.py_name }}Base(ServiceBase):
+class {{ (service.py_name + "Base") | add_to_all }}(ServiceBase):
     {% if service.comment %}
     """
     {{ service.comment | indent(4) }}

From 65eb7d28975e2de65624a1bd559cbc7125d4aa04 Mon Sep 17 00:00:00 2001
From: Adrien <adrien.vannson@protonmail.com>
Date: Fri, 2 May 2025 20:36:59 +0200
Subject: [PATCH 2/2] Remove comment

---
 src/betterproto2_compiler/templates/header.py.j2 | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/src/betterproto2_compiler/templates/header.py.j2 b/src/betterproto2_compiler/templates/header.py.j2
index 74738fcc..64a2de57 100644
--- a/src/betterproto2_compiler/templates/header.py.j2
+++ b/src/betterproto2_compiler/templates/header.py.j2
@@ -5,19 +5,6 @@
 # plugin: python-betterproto2
 # This file has been @generated
 
-{# __all__ = (
-    {% for _, enum in output_file.enums|dictsort(by="key") %}
-        "{{ enum.py_name }}",
-    {%- endfor -%}
-    {% for _, message in output_file.messages|dictsort(by="key") %}
-        "{{ message.py_name }}",
-    {%- endfor -%}
-    {% for _, service in output_file.services|dictsort(by="key") %}
-        "{{ service.py_name }}Stub",
-        "{{ service.py_name }}Base",
-    {%- endfor -%}
-) #}
-
 __all__ = (
     {%- for name in all -%}
     "{{ name }}",