From 1a23f09c16e985401a7b494df975d478c054feb3 Mon Sep 17 00:00:00 2001
From: Adrien <adrien.vannson@protonmail.com>
Date: Thu, 7 Nov 2024 04:11:45 +0100
Subject: [PATCH] Define __all__ (#625)

* Define __all__

* Use tuple instead of list

* Add test
---
 src/betterproto/templates/header.py.j2 | 14 ++++++++++++++
 tests/test_all_definition.py           | 19 +++++++++++++++++++
 2 files changed, 33 insertions(+)
 create mode 100644 tests/test_all_definition.py

diff --git a/src/betterproto/templates/header.py.j2 b/src/betterproto/templates/header.py.j2
index 9c8dddd..b6d0a6c 100644
--- a/src/betterproto/templates/header.py.j2
+++ b/src/betterproto/templates/header.py.j2
@@ -2,6 +2,20 @@
 # sources: {{ ', '.join(output_file.input_filenames) }}
 # plugin: python-betterproto
 # This file has been @generated
+
+__all__ = (
+    {%- for enum in output_file.enums -%}
+        "{{ enum.py_name }}",
+    {%- endfor -%}
+    {%- for message in output_file.messages -%}
+        "{{ message.py_name }}",
+    {%- endfor -%}
+    {%- for service in output_file.services -%}
+        "{{ service.py_name }}Stub",
+        "{{ service.py_name }}Base",
+    {%- endfor -%}
+)
+
 {% for i in output_file.python_module_imports|sort %}
 import {{ i }}
 {% endfor %}
diff --git a/tests/test_all_definition.py b/tests/test_all_definition.py
new file mode 100644
index 0000000..61abb5f
--- /dev/null
+++ b/tests/test_all_definition.py
@@ -0,0 +1,19 @@
+def test_all_definition():
+    """
+    Check that a compiled module defines __all__ with the right value.
+
+    These modules have been chosen since they contain messages, services and enums.
+    """
+    import tests.output_betterproto.enum as enum
+    import tests.output_betterproto.service as service
+
+    assert service.__all__ == (
+        "ThingType",
+        "DoThingRequest",
+        "DoThingResponse",
+        "GetThingRequest",
+        "GetThingResponse",
+        "TestStub",
+        "TestBase",
+    )
+    assert enum.__all__ == ("Choice", "ArithmeticOperator", "Test")