Add enums to the default __init__.py (#271)

software-dov · web-flow · commit 51b686132226 · 2020-01-22T13:48:23.000-08:00
Fix for #265
diff --git a/packages/gapic-generator/gapic/templates/%namespace/%name/__init__.py.j2 b/packages/gapic-generator/gapic/templates/%namespace/%name/__init__.py.j2
@@ -2,6 +2,7 @@
 
 {% block content %}
 {#  Import subpackages. -#}
+{% filter sort_lines -%}
 {% for subpackage in api.subpackages.keys() -%}
 from {% if api.naming.module_namespace %}{{ api.naming.module_namespace|join('.') }}.{% endif -%}
     {{ api.naming.versioned_module_name }} import {{ subpackage }}
@@ -14,34 +15,51 @@ from {% if api.naming.module_namespace %}{{ api.naming.module_namespace|join('.'
     {{ api.naming.versioned_module_name }}.services.{{ service.name|snake_case }}.client import {{ service.client_name }}
 {% endfor -%}
 
-{#  Import messages from each proto.
+{#  Import messages and enums from each proto.
     It is safe to import all of the messages into the same namespace here,
     because protocol buffers itself enforces selector uniqueness within
     a proto package.
 -#}
+{#  Import messages from each proto.
+    It is safe to import all of the messages into the same namespace here,
+    because protocol buffers itself enforces selector uniqueness within
+    a proto package.
+  -#}
 {% for proto in api.protos.values()|sort(attribute='module_name')
         if proto.meta.address.subpackage == api.subpackage_view -%}
-{% for message in proto.messages.values()|sort(attribute='name') -%}
+  {% for message in proto.messages.values()|sort(attribute='name') -%}
 from {% if api.naming.module_namespace %}{{ api.naming.module_namespace|join('.') }}.{% endif -%}
     {{ api.naming.versioned_module_name }}.types.{{ proto.module_name }} import {{ message.name }}
-{% endfor %}{% endfor %}
-
+{% endfor -%}
+{% for enum in proto.enums.values()|sort(attribute='name') -%}
+from {% if api.naming.module_namespace %}{{ api.naming.module_namespace|join('.') }}.{% endif -%}
+    {{ api.naming.versioned_module_name }}.types.{{ proto.module_name }} import {{ enum.name }}
+{% endfor %}{% endfor -%}
+{% endfilter %}
 {#  Define __all__.
     This requires the full set of imported names, so we iterate over
     them again.
 -#}
 __all__ = (
-    {%- for subpackage in api.subpackages.keys() %}
-    '{{ subpackage }}',
-    {%- endfor %}
-    {%- for service in api.services.values()|sort(attribute='name')
-            if service.meta.address.subpackage == api.subpackage_view %}
-    '{{ service.client_name }}',
-    {%- endfor %}
-    {%- for proto in api.protos.values()|sort(attribute='module_name')
-            if proto.meta.address.subpackage == api.subpackage_view %}
-    {%- for message in proto.messages.values()|sort(attribute='name') %}
-    '{{ message.name }}',
-    {%- endfor %}{% endfor %}
+{%- filter indent %}
+{% filter sort_lines -%}
+{% for subpackage in api.subpackages.keys() -%}
+'{{ subpackage }}',
+{% endfor -%}
+{% for service in api.services.values()|sort(attribute='name')
+   if service.meta.address.subpackage == api.subpackage_view -%}
+'{{ service.client_name }}',
+{% endfor -%}
+{% for proto in api.protos.values()|sort(attribute='module_name')
+   if proto.meta.address.subpackage == api.subpackage_view -%}
+{% for message in proto.messages.values()|sort(attribute='name') -%}
+'{{ message.name }}',
+{% endfor -%}
+{% for enum in proto.enums.values()|sort(attribute='name')
+   if proto.meta.address.subpackage == api.subpackage_view -%}
+'{{ enum.name }}',
+{% endfor -%}{% endfor -%}
+{% endfilter -%}
+{% endfilter -%}
 )
 {% endblock %}
diff --git a/packages/gapic-generator/gapic/utils/lines.py b/packages/gapic-generator/gapic/utils/lines.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import textwrap
+from typing import Iterable
 
 
 def sort_lines(text: str, dedupe: bool = True) -> str:
@@ -27,11 +28,11 @@ def sort_lines(text: str, dedupe: bool = True) -> str:
     trailing = '\n' if text.endswith('\n') else ''
 
     # Split the text into individual lines, throwing away any empty lines.
-    lines = [i for i in text.strip().split('\n') if i.strip()]
+    lines: Iterable[str] = (i for i in text.strip().split('\n') if i.strip())
 
     # De-duplicate the lines if requested.
     if dedupe:
-        lines = list(set(lines))
+        lines = set(lines)
 
     # Return the final string.
     answer = '\n'.join(sorted(lines))
@@ -78,8 +79,8 @@ def wrap(text: str, width: int, *, offset: int = None, indent: int = 0) -> str:
     first = text.split('\n')[0] + '\n'
     if len(first) > width - offset:
         initial = textwrap.wrap(first,
-            break_long_words=False,
-            width=width - offset,
+                                break_long_words=False,
+                                width=width - offset,
                                 )
         # Strip the first \n from the text so it is not misidentified as an
         # intentionally short line below.
