perf: reduce unnecessary copies, optimize Address comparison (#855)

Includes various other performance optimizations and minor whitespace fixes.

Local testing cuts down Ads generation time to ~1m (and prevents OOM) and DialogflowCX from ~10 minutes to ~5 minutes.

Author: software-dov

gapic/ads-templates/%namespace/%name/%version/%sub/__init__.py.j2

@@ -8,21 +8,23 @@
     them again.
 -#}
 __all__ = (
-    {% for subpackage in api.subpackages|dictsort %}
+    {% filter sort_lines -%}
+    {% for subpackage in api.subpackages -%}
     '{{ subpackage }}',
-    {% endfor %}
-    {% for service in api.services.values()|sort(attribute='client_name')
-            if service.meta.address.subpackage == api.subpackage_view %}
+    {% endfor -%}
+    {% for service in api.services.values()
+            if service.meta.address.subpackage == api.subpackage_view -%}
     '{{ service.client_name }}',
-    {% endfor %}
-    {% for proto in api.protos.values()|sort(attribute='name')
-            if proto.meta.address.subpackage == api.subpackage_view %}
-    {% for message in proto.messages.values() %}
+    {% endfor -%}
+    {% for proto in api.protos.values()
+            if proto.meta.address.subpackage == api.subpackage_view -%}
+    {% for message in proto.messages.values() -%}
     '{{ message.name }}',
-    {% endfor %}
-    {% for enum in proto.enums.values()|sort(attribute='name') %}
+    {% endfor -%}
+    {% for enum in proto.enums.values() -%}
     '{{ enum.name }}',
-    {% endfor %}
-    {% endfor %}
+    {% endfor -%}
+    {% endfor -%}
+    {% endfilter -%}
 )
 {% endblock %}

gapic/ads-templates/%namespace/%name/%version/__init__.py.j2

@@ -1,7 +1,7 @@
 {% extends '_base.py.j2' %}
 {% block content %}
 
-{% if opts.lazy_import %} {# lazy import #}
+{% if opts.lazy_import -%} {# lazy import #}
 import importlib
 import sys
 

gapic/ads-templates/%namespace/%name/__init__.py.j2

@@ -1,7 +1,7 @@
 {% extends '_base.py.j2' %}
 {% block content %}
 
-{% if opts.lazy_import %} {# lazy import #}
+{% if opts.lazy_import -%} {# lazy import #}
 import importlib
 import sys
 

gapic/ads-templates/tests/unit/gapic/%name_%version/test_module_import.py.j2

@@ -1,7 +1,7 @@
 {% extends "_base.py.j2" %}
 {% block content %}
 
-{% if opts.lazy_import %} {# lazy import #}
+{% if opts.lazy_import -%} {# lazy import #}
 import pytest
 
 

gapic/generator/generator.py

@@ -378,7 +378,7 @@ def _get_filename(
         # Replace the %namespace variable.
         filename = filename.replace(
             "%namespace",
-            os.path.sep.join([i.lower() for i in api_schema.naming.namespace]),
+            os.path.sep.join(i.lower() for i in api_schema.naming.namespace),
         ).lstrip(os.path.sep)
 
         # Replace the %name, %version, and %sub variables.

gapic/schema/metadata.py

@@ -36,22 +36,30 @@
 from gapic.utils import cached_property
 from gapic.utils import RESERVED_NAMES
 
+# This class is a minor hack to optimize Address's __eq__ method.
+
 
 @dataclasses.dataclass(frozen=True)
-class Address:
+class BaseAddress:
     name: str = ''
     module: str = ''
     module_path: Tuple[int, ...] = dataclasses.field(default_factory=tuple)
     package: Tuple[str, ...] = dataclasses.field(default_factory=tuple)
     parent: Tuple[str, ...] = dataclasses.field(default_factory=tuple)
+
+
+@dataclasses.dataclass(frozen=True)
+class Address(BaseAddress):
     api_naming: naming.Naming = dataclasses.field(
         default_factory=naming.NewNaming,
     )
     collisions: FrozenSet[str] = dataclasses.field(default_factory=frozenset)
 
     def __eq__(self, other) -> bool:
-        return all(getattr(self, i) == getattr(other, i)
-                   for i in ('name', 'module', 'module_path', 'package', 'parent'))
+        # We don't want to use api_naming or collisions to determine equality,
+        # so defer to the parent class's eq method.
+        # This is an fairly important optimization for large APIs.
+        return super().__eq__(other)
 
     def __hash__(self):
         # Do NOT include collisions; they are not relevant.
@@ -94,7 +102,8 @@ def __str__(self) -> str:
         # Return the Python identifier.
         return '.'.join(self.parent + (self.name,))
 
-    def __repr__(self) -> str:
+    @cached_property
+    def __cached_string_repr(self):
         return "({})".format(
             ", ".join(
                 (
@@ -108,6 +117,9 @@ def __repr__(self) -> str:
             )
         )
 
+    def __repr__(self) -> str:
+        return self.__cached_string_repr
+
     @property
     def module_alias(self) -> str:
         """Return an appropriate module alias if necessary.
@@ -118,16 +130,15 @@ def module_alias(self) -> str:
         while still providing names that are fundamentally readable
         to users (albeit looking auto-generated).
         """
-        if self.module in self.collisions | RESERVED_NAMES:
+        # This is a minor optimization to prevent constructing a temporary set.
+        if self.module in self.collisions or self.module in RESERVED_NAMES:
             return '_'.join(
                 (
                     ''.join(
-                        [
-                            partial_name[0]
-                            for i in self.package
-                            for partial_name in i.split("_")
-                            if i != self.api_naming.version
-                        ]
+                        partial_name[0]
+                        for i in self.package
+                        for partial_name in i.split("_")
+                        if i != self.api_naming.version
                     ),
                     self.module,
                 )
@@ -302,7 +313,11 @@ def with_context(self, *, collisions: FrozenSet[str]) -> 'Address':
         ``Address`` object aliases module names to avoid naming collisions in
         the file being written.
         """
-        return dataclasses.replace(self, collisions=collisions)
+        return (
+            dataclasses.replace(self, collisions=collisions)
+            if collisions and collisions != self.collisions
+            else self
+        )
 
 
 @dataclasses.dataclass(frozen=True)
@@ -340,7 +355,7 @@ def with_context(self, *, collisions: FrozenSet[str]) -> 'Metadata':
         return dataclasses.replace(
             self,
             address=self.address.with_context(collisions=collisions),
-        )
+        ) if collisions and collisions != self.address.collisions else self
 
 
 @dataclasses.dataclass(frozen=True)

gapic/schema/wrappers.py

@@ -463,24 +463,24 @@ def with_context(self, *,
         visited_messages = visited_messages | {self}
         return dataclasses.replace(
             self,
-            fields=collections.OrderedDict(
-                (k, v.with_context(
+            fields={
+                k: v.with_context(
                     collisions=collisions,
                     visited_messages=visited_messages
-                ))
-                for k, v in self.fields.items()
-            ) if not skip_fields else self.fields,
-            nested_enums=collections.OrderedDict(
-                (k, v.with_context(collisions=collisions))
+                ) for k, v in self.fields.items()
+            } if not skip_fields else self.fields,
+            nested_enums={
+                k: v.with_context(collisions=collisions)
                 for k, v in self.nested_enums.items()
-            ),
-            nested_messages=collections.OrderedDict(
-                (k, v.with_context(
+            },
+            nested_messages={
+                k: v.with_context(
                     collisions=collisions,
                     skip_fields=skip_fields,
                     visited_messages=visited_messages,
-                ))
-                for k, v in self.nested_messages.items()),
+                )
+                for k, v in self.nested_messages.items()
+            },
             meta=self.meta.with_context(collisions=collisions),
         )
 
@@ -535,7 +535,7 @@ def with_context(self, *, collisions: FrozenSet[str]) -> 'EnumType':
         return dataclasses.replace(
             self,
             meta=self.meta.with_context(collisions=collisions),
-        )
+        ) if collisions else self
 
     @property
     def options_dict(self) -> Dict:
@@ -957,9 +957,11 @@ def with_context(self, *, collisions: FrozenSet[str]) -> 'Method':
         ``Method`` object aliases module names to avoid naming collisions
         in the file being written.
         """
-        maybe_lro = self.lro.with_context(
-            collisions=collisions
-        ) if self.lro else None
+        maybe_lro = None
+        if self.lro:
+            maybe_lro = self.lro.with_context(
+                collisions=collisions
+            ) if collisions else self.lro
 
         return dataclasses.replace(
             self,
@@ -1195,13 +1197,12 @@ def with_context(self, *, collisions: FrozenSet[str]) -> 'Service':
         """
         return dataclasses.replace(
             self,
-            methods=collections.OrderedDict(
-                (k, v.with_context(
-                    # A methodd's flattened fields create additional names
+            methods={
+                k: v.with_context(
+                    # A method's flattened fields create additional names
                     # that may conflict with module imports.
                     collisions=collisions | frozenset(v.flattened_fields.keys()))
-                 )
                 for k, v in self.methods.items()
-            ),
+            },
             meta=self.meta.with_context(collisions=collisions),
         )

gapic/templates/%namespace/%name_%version/%sub/__init__.py.j2

@@ -33,21 +33,23 @@ from .types.{{ proto.module_name }} import {{ enum.name }}
     them again.
 -#}
 __all__ = (
-    {% for subpackage in api.subpackages|dictsort %}
+    {% filter sort_lines -%}
+    {% for subpackage in api.subpackages -%}
     '{{ subpackage }}',
-    {% endfor %}
-    {% for service in api.services.values()|sort(attribute='client_name')
-            if service.meta.address.subpackage == api.subpackage_view %}
+    {% endfor -%}
+    {% for service in api.services.values()
+            if service.meta.address.subpackage == api.subpackage_view -%}
     '{{ service.client_name }}',
-    {% endfor %}
-    {% for proto in api.protos.values()|sort(attribute='name')
-            if proto.meta.address.subpackage == api.subpackage_view %}
-    {% for message in proto.messages.values()|sort(attribute='name') %}
+    {% endfor -%}
+    {% for proto in api.protos.values()
+            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') %}
+    {% endfor -%}
+    {% for enum in proto.enums.values() -%}
     '{{ enum.name }}',
-    {% endfor %}
-    {% endfor %}
+    {% endfor -%}
+    {% endfor -%}
+    {% endfilter %}
 )
 {% endblock %}

tests/unit/schema/test_metadata.py

@@ -33,7 +33,7 @@ def test_address_str_with_context():
         package=('foo', 'bar'),
         module='baz',
         name='Bacon',
-    ).with_context(collisions={'baz'})
+    ).with_context(collisions=frozenset({'baz'}))
     assert str(addr) == 'fb_baz.Bacon'
 
 

tests/unit/schema/wrappers/test_message.py

@@ -56,7 +56,7 @@ def test_message_ident():
 
 def test_message_ident_collisions():
     message = make_message('Baz', package='foo.v1', module='bar').with_context(
-        collisions={'bar'},
+        collisions=frozenset({'bar'}),
     )
     assert str(message.ident) == 'fv_bar.Baz'
     assert message.ident.sphinx == 'foo.v1.bar.Baz'