Define __all__ (#625)

AdrienVannson · web-flow · commit 1a23f09c16e9 · 2024-11-07T03:11:45.000Z
* Define __all__

* Use tuple instead of list

* Add test
diff --git 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
@@ -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")
