Fix for flattened unit tests with deeply nested request messages (#346)

Expands the unit-test visible method ref types to include _all_
recursive types.

Do NOT use this deep nesting in the client submodule. The client ONLY
needs to import, for each method, the request type, response type, and
the types of any flattened fields.

Author: software-dov

packages/gapic-generator/gapic/schema/api.py

@@ -19,10 +19,8 @@
 
 import collections
 import dataclasses
-import keyword
 import os
 import sys
-from itertools import chain
 from typing import Callable, Container, Dict, FrozenSet, Mapping, Optional, Sequence, Set, Tuple
 
 from google.api_core import exceptions  # type: ignore
@@ -37,6 +35,7 @@
 from gapic.schema import naming as api_naming
 from gapic.utils import cached_property
 from gapic.utils import to_snake_case
+from gapic.utils import RESERVED_NAMES
 
 
 @dataclasses.dataclass(frozen=True)
@@ -130,13 +129,13 @@ def names(self) -> FrozenSet[str]:
         # from distinct packages.
         modules: Dict[str, Set[str]] = collections.defaultdict(set)
         for m in self.all_messages.values():
-            for t in m.field_types:
+            for t in m.recursive_field_types:
                 modules[t.ident.module].add(t.ident.package)
 
         answer.update(
             module_name
             for module_name, packages in modules.items()
-            if len(packages) > 1
+            if len(packages) > 1 or module_name in RESERVED_NAMES
         )
 
         # Return the set of collision names.
@@ -229,7 +228,7 @@ def disambiguate_keyword_fname(
                 visited_names: Container[str]) -> str:
             path, fname = os.path.split(full_path)
             name, ext = os.path.splitext(fname)
-            if name in keyword.kwlist or full_path in visited_names:
+            if name in RESERVED_NAMES or full_path in visited_names:
                 name += "_"
                 full_path = os.path.join(path, name + ext)
                 if full_path in visited_names:

packages/gapic-generator/gapic/schema/wrappers.py

@@ -31,7 +31,7 @@
 import dataclasses
 import re
 from itertools import chain
-from typing import (cast, Dict, FrozenSet, List, Mapping, Optional,
+from typing import (cast, Dict, FrozenSet, Iterable, List, Mapping, Optional,
                     Sequence, Set, Union)
 
 from google.api import annotations_pb2      # type: ignore
@@ -225,7 +225,6 @@ def __hash__(self):
 
     @utils.cached_property
     def field_types(self) -> Sequence[Union['MessageType', 'EnumType']]:
-        """Return all composite fields used in this proto's messages."""
         answer = tuple(
             field.type
             for field in self.fields.values()
@@ -234,6 +233,23 @@ def field_types(self) -> Sequence[Union['MessageType', 'EnumType']]:
 
         return answer
 
+    @utils.cached_property
+    def recursive_field_types(self) -> Sequence[
+        Union['MessageType', 'EnumType']
+    ]:
+        """Return all composite fields used in this proto's messages."""
+        types: List[Union['MessageType', 'EnumType']] = []
+        stack = [iter(self.fields.values())]
+        while stack:
+            fields_iter = stack.pop()
+            for field in fields_iter:
+                if field.message and field.type not in types:
+                    stack.append(iter(field.message.fields.values()))
+                if not field.is_primitive:
+                    types.append(field.type)
+
+        return tuple(types)
+
     @property
     def map(self) -> bool:
         """Return True if the given message is a map, False otherwise."""
@@ -654,19 +670,30 @@ def paged_result_field(self) -> Optional[Field]:
 
     @utils.cached_property
     def ref_types(self) -> Sequence[Union[MessageType, EnumType]]:
+        return self._ref_types(True)
+
+    @utils.cached_property
+    def flat_ref_types(self) -> Sequence[Union[MessageType, EnumType]]:
+        return self._ref_types(False)
+
+    def _ref_types(self, recursive: bool) -> Sequence[Union[MessageType, EnumType]]:
         """Return types referenced by this method."""
         # Begin with the input (request) and output (response) messages.
-        answer = [self.input]
+        answer: List[Union[MessageType, EnumType]] = [self.input]
+        types: Iterable[Union[MessageType, EnumType]] = (
+            self.input.recursive_field_types if recursive
+            else (
+                f.type
+                for f in self.flattened_fields.values()
+                if f.message or f.enum
+            )
+        )
+        answer.extend(types)
+
         if not self.void:
             answer.append(self.client_output)
             answer.extend(self.client_output.field_types)
 
-        answer.extend(
-            field.type
-            for field in self.flattened_fields.values()
-            if field.message or field.enum
-        )
-
         # If this method has LRO, it is possible (albeit unlikely) that
         # the LRO messages reside in a different module.
         if self.lro:

packages/gapic-generator/gapic/templates/%namespace/%name_%version/%sub/services/%service/client.py.j2

@@ -14,7 +14,7 @@ from google.oauth2 import service_account              # type: ignore
 
 {% filter sort_lines -%}
 {% for method in service.methods.values() -%}
-{% for ref_type in method.ref_types -%}
+{% for ref_type in method.flat_ref_types -%}
 {{ ref_type.ident.python_import }}
 {% endfor -%}
 {% endfor -%}

packages/gapic-generator/gapic/utils/__init__.py

@@ -21,6 +21,7 @@
 from gapic.utils.filename import to_valid_module_name
 from gapic.utils.lines import sort_lines
 from gapic.utils.lines import wrap
+from gapic.utils.reserved_names import RESERVED_NAMES
 from gapic.utils.rst import rst
 
 
@@ -29,6 +30,7 @@
     'doc',
     'empty',
     'partition',
+    'RESERVED_NAMES',
     'rst',
     'sort_lines',
     'to_snake_case',

packages/gapic-generator/gapic/utils/reserved_names.py

@@ -0,0 +1,18 @@
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import keyword
+
+
+RESERVED_NAMES = frozenset(keyword.kwlist)

packages/gapic-generator/test_utils/test_utils.py

@@ -145,13 +145,15 @@ def get_message(dot_path: str, *,
 
 
 def make_method(
-        name: str, input_message: wrappers.MessageType = None,
+        name: str,
+        input_message: wrappers.MessageType = None,
         output_message: wrappers.MessageType = None,
         package: typing.Union[typing.Tuple[str], str] = 'foo.bar.v1',
         module: str = 'baz',
         http_rule: http_pb2.HttpRule = None,
         signatures: typing.Sequence[str] = (),
-        **kwargs) -> wrappers.Method:
+        **kwargs
+) -> wrappers.Method:
     # Use default input and output messages if they are not provided.
     input_message = input_message or make_message('MethodInput')
     output_message = output_message or make_message('MethodOutput')
@@ -229,11 +231,14 @@ def make_field(
     )
 
 
-def make_message(name: str, package: str = 'foo.bar.v1', module: str = 'baz',
-                 fields: typing.Sequence[wrappers.Field] = (),
-                 meta: metadata.Metadata = None,
-                 options: desc.MethodOptions = None,
-                 ) -> wrappers.MessageType:
+def make_message(
+    name: str,
+    package: str = 'foo.bar.v1',
+    module: str = 'baz',
+    fields: typing.Sequence[wrappers.Field] = (),
+    meta: metadata.Metadata = None,
+    options: desc.MethodOptions = None,
+) -> wrappers.MessageType:
     message_pb = desc.DescriptorProto(
         name=name,
         field=[i.field_pb for i in fields],

packages/gapic-generator/tests/unit/schema/wrappers/test_message.py

@@ -70,7 +70,15 @@ def test_get_field():
 
 def test_field_types():
     # Create the inner message.
-    inner_msg = make_message('InnerMessage', fields=())
+    inner_msg = make_message(
+        'InnerMessage',
+        fields=(
+            make_field(
+                'hidden_message',
+                message=make_message('HiddenMessage'),
+            ),
+        )
+    )
     inner_enum = make_enum('InnerEnum')
 
     # Create the outer message, which contains an Inner as a field.
@@ -87,6 +95,33 @@ def test_field_types():
     assert inner_enum in outer.field_types
 
 
+def test_field_types_recursive():
+    enumeration = make_enum('Enumeration')
+    innest_msg = make_message(
+        'InnestMessage',
+        fields=(
+            make_field('enumeration', enum=enumeration),
+        )
+    )
+    inner_msg = make_message(
+        'InnerMessage',
+        fields=(
+            make_field('innest_message', message=innest_msg),
+        )
+    )
+    topmost_msg = make_message(
+        'TopmostMessage',
+        fields=(
+            make_field('inner_message', message=inner_msg),
+            make_field('uninteresting')
+        )
+    )
+
+    actual = {t.name for t in topmost_msg.recursive_field_types}
+    expected = {t.name for t in (enumeration, innest_msg, inner_msg)}
+    assert actual == expected
+
+
 def test_get_field_recursive():
     # Create the inner message.
     inner_fields = (make_field('zero'), make_field('one'))

packages/gapic-generator/tests/unit/schema/wrappers/test_method.py

@@ -23,6 +23,7 @@
 from gapic.schema import wrappers
 
 from test_utils.test_utils import (
+    make_enum,
     make_field,
     make_message,
     make_method,
@@ -151,6 +152,53 @@ def test_method_paged_result_ref_types():
     }
 
 
+def test_flattened_ref_types():
+    method = make_method(
+        'IdentifyMollusc',
+        input_message=make_message(
+            'IdentifyMolluscRequest',
+            fields=(
+                make_field(
+                    'cephalopod',
+                    message=make_message(
+                        'Cephalopod',
+                        fields=(
+                            make_field('mass_kg', type='TYPE_INT32'),
+                            make_field(
+                                'squid',
+                                number=2,
+                                message=make_message('Squid'),
+                            ),
+                            make_field(
+                                'clam',
+                                number=3,
+                                message=make_message('Clam'),
+                            ),
+                        ),
+                    ),
+                ),
+                make_field(
+                    'stratum',
+                    enum=make_enum(
+                        'Stratum',
+                    )
+                ),
+            ),
+        ),
+        signatures=('cephalopod.squid,stratum',),
+        output_message=make_message('Mollusc'),
+    )
+
+    expected_flat_ref_type_names = {
+        'IdentifyMolluscRequest',
+        'Squid',
+        'Stratum',
+        'Mollusc',
+    }
+    actual_flat_ref_type_names = {t.name for t in method.flat_ref_types}
+    assert expected_flat_ref_type_names == actual_flat_ref_type_names
+
+
 def test_method_field_headers_none():
     method = make_method('DoSomething')
     assert isinstance(method.field_headers, collections.abc.Sequence)