diff --git a/src/betterproto/templates/header.py.j2 b/src/betterproto/templates/header.py.j2 index 9c8ddddc..b6d0a6c4 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 00000000..61abb5f3 --- /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")