Correctly find the sample template when run via protoc (#198)

Add a sample_template parameter to sample generation
Adjust tests to pass the parameter

Author: software-dov

packages/gapic-generator/gapic/generator/generator.py

@@ -97,14 +97,18 @@ def get_response(self, api_schema: api.API) -> CodeGeneratorResponse:
                                                       api_schema=api_schema,
                                                       ))
 
-        output_files.update(self._generate_samples_and_manifest(api_schema))
+        output_files.update(self._generate_samples_and_manifest(
+            api_schema,
+            self._env.get_template(sample_templates[0]),
+        ))
 
         # Return the CodeGeneratorResponse output.
         return CodeGeneratorResponse(file=[i for i in output_files.values()])
 
     def _generate_samples_and_manifest(
-        self,
-        api_schema: api.API
+            self,
+            api_schema: api.API,
+            sample_template: jinja2.Template,
     ) -> Dict[str, CodeGeneratorResponse.File]:
         """Generate samples and samplegen manifest for the API.
 
@@ -152,7 +156,11 @@ def _generate_samples_and_manifest(
                         str(spec).encode('utf8')).hexdigest()[:8]
                     spec["id"] += f"_{spec_hash}"
 
-                sample = samplegen.generate_sample(spec, self._env, api_schema)
+                sample = samplegen.generate_sample(
+                    spec,
+                    api_schema,
+                    sample_template,
+                )
 
                 fpath = spec["id"] + ".py"
                 fpath_to_spec_and_rendered[os.path.join(out_dir, fpath)] = (spec,

packages/gapic-generator/gapic/samplegen/manifest.py

@@ -84,15 +84,15 @@ def transform_path(fpath):
             yaml.Collection(
                 name="samples",
                 elements=[
-                    [  # type: ignore
+                    [
                         # Mypy doesn't correctly intuit the type of the
                         # "region_tag" conditional expression.
                         yaml.Alias(environment.anchor_name or ""),
                         yaml.KeyVal("sample", sample["id"]),
                         yaml.KeyVal(
                             "path", transform_path(fpath)
                         ),
-                        (yaml.KeyVal("region_tag", sample["region_tag"])
+                        (yaml.KeyVal("region_tag", sample["region_tag"])  # type: ignore
                          if "region_tag" in sample else
                          yaml.Null),
                     ]

packages/gapic-generator/gapic/samplegen/samplegen.py

@@ -21,7 +21,7 @@
 import time
 
 from gapic.samplegen_utils import types
-from gapic.schema import (api, wrappers)
+from gapic.schema import wrappers
 
 from collections import (defaultdict, namedtuple, ChainMap as chainmap)
 from typing import (ChainMap, Dict, List, Mapping, Optional, Tuple)
@@ -139,8 +139,8 @@ def __init__(self, method: wrappers.Method):
             # and whether it's an enum or a message or a primitive type.
             # The method call response isn't a field, so construct an artificial
             # field that wraps the response.
-            {  # type: ignore
-                "$resp": MockField(response_type, False)
+            {
+                "$resp": MockField(response_type, False)  # type: ignore
             }
         )
 
@@ -481,8 +481,10 @@ def _validate_format(self, body: List[str]):
         num_prints = fmt_str.count("%s")
         if num_prints != len(body) - 1:
             raise types.MismatchedFormatSpecifier(
-                "Expected {} expresssions in format string but received {}".format(
-                    num_prints, len(body) - 1
+                "Expected {} expresssions in format string '{}' but found {}".format(
+                    num_prints,
+                    fmt_str,
+                    len(body) - 1
                 )
             )
 
@@ -502,7 +504,7 @@ def _validate_define(self, body: str):
         """
         # Note: really checking for safety would be equivalent to
         #       re-implementing the python interpreter.
-        m = re.match(r"^([a-zA-Z]\w*)=([^=]+)$", body)
+        m = re.match(r"^([a-zA-Z_]\w*) *= *([^=]+)$", body)
         if not m:
             raise types.BadAssignment(f"Bad assignment statement: {body}")
 
@@ -654,27 +656,23 @@ def _validate_loop(self, loop):
     }
 
 
-def generate_sample(sample,
-                    env: jinja2.environment.Environment,
-                    api_schema: api.API,
-                    template_name: str = DEFAULT_TEMPLATE_NAME) -> str:
+def generate_sample(
+        sample,
+        api_schema,
+        sample_template: jinja2.Template
+) -> str:
     """Generate a standalone, runnable sample.
 
     Rendering and writing the rendered output is left for the caller.
 
     Args:
         sample (Any): A definition for a single sample generated from parsed yaml.
-        env (jinja2.environment.Environment): The jinja environment used to generate
-                                              the filled template for the sample.
         api_schema (api.API): The schema that defines the API to which the sample belongs.
-        template_name (str): An optional override for the name of the template
-                             used to generate the sample.
+        sample_template (jinja2.Template): The template representing a generic sample.
 
     Returns:
         str: The rendered sample.
     """
-    sample_template = env.get_template(template_name)
-
     service_name = sample["service"]
     service = api_schema.services.get(service_name)
     if not service:
@@ -701,7 +699,12 @@ def generate_sample(sample,
     return sample_template.render(
         file_header=FILE_HEADER,
         sample=sample,
-        imports=[],
+        imports=[
+            "from google import auth",
+            "from google.auth import credentials",
+        ],
         calling_form=calling_form,
         calling_form_enum=types.CallingForm,
+        api=api_schema,
+        service=service,
     )

packages/gapic-generator/gapic/templates/examples/feature_fragments.j2

@@ -191,32 +191,32 @@ client.{{ sample.rpc|snake_case }}({{ render_request_params(sample.request) }})
 {# it's just easier to set up client side streaming and other things from outside this macro. #}
 {% macro render_calling_form(method_invocation_text, calling_form, calling_form_enum, response_statements ) %}
 {% if calling_form == calling_form_enum.Request %}
-response = {{ method_invocation_text }}
+response = {{ method_invocation_text|trim }}
 {% for statement in response_statements %}
 {{ dispatch_statement(statement)|trim }}
 {% endfor %}
 {% elif calling_form == calling_form_enum.RequestPagedAll %}
-page_result = {{ method_invocation_text }}
+page_result = {{ method_invocation_text|trim }}
 for response in page_result:
   {% for statement in response_statements %}
     {{ dispatch_statement(statement)|trim }}
   {% endfor %}
 {% elif calling_form == calling_form_enum.RequestPaged %}
-page_result = {{ method_invocation_text}}
+page_result = {{ method_invocation_text|trim }}
 for page in page_result.pages():
     for response in page:
   {% for statement in response_statements %}
         {{ dispatch_statement(statement)|trim }}
   {% endfor %}
 {% elif calling_form in [calling_form_enum.RequestStreamingServer,
                         calling_form_enum.RequestStreamingBidi] %}
-stream = {{ method_invocation_text }}
+stream = {{ method_invocation_text|trim }}
 for response in stream:
   {% for statement in response_statements %}
     {{ dispatch_statement(statement)|trim }}
   {% endfor %}
 {% elif calling_form == calling_form_enum.LongRunningRequestPromise %}
-operation = {{ method_invocation_text }}
+operation = {{ method_invocation_text|trim }}
 
 print("Waiting for operation to complete...")
 
@@ -237,8 +237,8 @@ def main():
 
     parser = argparse.ArgumentParser()
 {% with arg_list = [] %}
-  {% for request in request_block if request.body -%}
-      {% for attr in request.body if attr.input_parameter %}
+{% for request in request_block if request.body -%}
+{% for attr in request.body if attr.input_parameter %}
     parser.add_argument("--{{ attr.input_parameter }}",
                         type=str,
                         default={{ attr.value }})

packages/gapic-generator/gapic/templates/examples/sample.py.j2

@@ -20,22 +20,24 @@
 {#                   callingFormEnum #}
 {# Note: this sample template is WILDLY INACCURATE AND INCOMPLETE #}
 {# It does not correctly enums, unions, top level attributes, or various other things #}
-{% import "feature_fragments.j2" as frags %}
+{% import "examples/feature_fragments.j2" as frags %}
 {{ frags.sample_header(file_header, sample, calling_form) }}
 
 # [START {{ sample.id }}]
 {# python code is responsible for all transformations: all we do here is render #}
 {% for import_statement in imports %}
 {{ import_statement }}
 {% endfor %}
+from {{ (api.naming.module_namespace + (api.naming.versioned_module_name,) + service.meta.address.subpackage)|join(".") }}.services.{{ service.name|snake_case }} import {{ service.name }}
 
 {# also need calling form #}
 def sample_{{ frags.render_method_name(sample.rpc)|trim -}}({{ frags.print_input_params(sample.request)|trim -}}):
     """{{ sample.description }}"""
 
-    client = {{ sample.service.split(".")[-3:-1]|
-                map("lower")|
-                join("_") }}.{{ sample.service.split(".")[-1] }}Client()
+    client = {{ service.name }}(
+        credentials=credentials.AnonymousCredentials(),
+        transport="grpc",
+    )
 
     {{ frags.render_request_setup(sample.request)|indent }}
 {% with method_call = frags.render_method_call(sample, calling_form, calling_form_enum) %}

packages/gapic-generator/tests/unit/generator/test_generator.py

@@ -42,12 +42,15 @@ def test_custom_template_directory():
 def test_get_response():
     g = make_generator()
     with mock.patch.object(jinja2.FileSystemLoader, 'list_templates') as lt:
-        lt.return_value = ['foo/bar/baz.py.j2']
+        lt.return_value = ['foo/bar/baz.py.j2', 'molluscs/squid/sample.py.j2']
         with mock.patch.object(jinja2.Environment, 'get_template') as gt:
             gt.return_value = jinja2.Template('I am a template result.')
             cgr = g.get_response(api_schema=make_api())
     lt.assert_called_once()
-    gt.assert_called_once()
+    gt.assert_has_calls([
+        mock.call('foo/bar/baz.py.j2'),
+        mock.call('molluscs/squid/sample.py.j2')
+    ])
     assert len(cgr.file) == 1
     assert cgr.file[0].name == 'foo/bar/baz.py'
     assert cgr.file[0].content == 'I am a template result.\n'
@@ -56,24 +59,34 @@ def test_get_response():
 def test_get_response_ignores_empty_files():
     g = make_generator()
     with mock.patch.object(jinja2.FileSystemLoader, 'list_templates') as lt:
-        lt.return_value = ['foo/bar/baz.py.j2']
+        lt.return_value = ['foo/bar/baz.py.j2', 'molluscs/squid/sample.py.j2']
         with mock.patch.object(jinja2.Environment, 'get_template') as gt:
             gt.return_value = jinja2.Template('# Meaningless comment')
             cgr = g.get_response(api_schema=make_api())
     lt.assert_called_once()
-    gt.assert_called_once()
+    gt.assert_has_calls([
+        mock.call('foo/bar/baz.py.j2'),
+        mock.call('molluscs/squid/sample.py.j2')
+    ])
     assert len(cgr.file) == 0
 
 
 def test_get_response_ignores_private_files():
     g = make_generator()
     with mock.patch.object(jinja2.FileSystemLoader, 'list_templates') as lt:
-        lt.return_value = ['foo/bar/baz.py.j2', 'foo/bar/_base.py.j2']
+        lt.return_value = [
+            'foo/bar/baz.py.j2',
+            'foo/bar/_base.py.j2',
+            'molluscs/squid/sample.py.j2',
+        ]
         with mock.patch.object(jinja2.Environment, 'get_template') as gt:
             gt.return_value = jinja2.Template('I am a template result.')
             cgr = g.get_response(api_schema=make_api())
     lt.assert_called_once()
-    gt.assert_called_once()
+    gt.assert_has_calls([
+        mock.call('foo/bar/baz.py.j2'),
+        mock.call('molluscs/squid/sample.py.j2')
+    ])
     assert len(cgr.file) == 1
     assert cgr.file[0].name == 'foo/bar/baz.py'
     assert cgr.file[0].content == 'I am a template result.\n'
@@ -82,7 +95,10 @@ def test_get_response_ignores_private_files():
 def test_get_response_fails_invalid_file_paths():
     g = make_generator()
     with mock.patch.object(jinja2.FileSystemLoader, 'list_templates') as lt:
-        lt.return_value = ['foo/bar/$service/$proto/baz.py.j2']
+        lt.return_value = [
+            'foo/bar/$service/$proto/baz.py.j2',
+            'molluscs/squid/sample.py.j2',
+        ]
         with pytest.raises(ValueError) as ex:
             g.get_response(api_schema=make_api())
 
@@ -93,7 +109,10 @@ def test_get_response_fails_invalid_file_paths():
 def test_get_response_enumerates_services():
     g = make_generator()
     with mock.patch.object(jinja2.FileSystemLoader, 'list_templates') as lt:
-        lt.return_value = ['foo/$service/baz.py.j2']
+        lt.return_value = [
+            'foo/$service/baz.py.j2',
+            'molluscs/squid/sample.py.j2',
+        ]
         with mock.patch.object(jinja2.Environment, 'get_template') as gt:
             gt.return_value = jinja2.Template('Service: {{ service.name }}')
             cgr = g.get_response(api_schema=make_api(make_proto(
@@ -112,7 +131,10 @@ def test_get_response_enumerates_services():
 def test_get_response_enumerates_proto():
     g = make_generator()
     with mock.patch.object(jinja2.FileSystemLoader, 'list_templates') as lt:
-        lt.return_value = ['foo/$proto.py.j2']
+        lt.return_value = [
+            'foo/$proto.py.j2',
+            'molluscs/squid/sample.py.j2',
+        ]
         with mock.patch.object(jinja2.Environment, 'get_template') as gt:
             gt.return_value = jinja2.Template('Proto: {{ proto.module_name }}')
             cgr = g.get_response(api_schema=make_api(
@@ -146,6 +168,7 @@ def test_get_response_divides_subpackages():
         lt.return_value = [
             'foo/$sub/types/$proto.py.j2',
             'foo/$sub/services/$service.py.j2',
+            'molluscs/squid/sample.py.j2',
         ]
         with mock.patch.object(jinja2.Environment, 'get_template') as gt:
             gt.return_value = jinja2.Template("""
@@ -274,7 +297,11 @@ def test_parse_sample_paths(fs):
 @mock.patch(
     'time.gmtime',
 )
-def test_samplegen_config_to_output_files(mock_gmtime, mock_generate_sample, fs):
+def test_samplegen_config_to_output_files(
+        mock_gmtime,
+        mock_generate_sample,
+        fs,
+):
     # These time values are nothing special,
     # they just need to be deterministic.
     returner = mock.MagicMock()
@@ -303,13 +330,14 @@ def test_samplegen_config_to_output_files(mock_gmtime, mock_generate_sample, fs)
         )
     )
 
-    mock_generate_sample
-
     g = generator.Generator(
         options.Options.build(
             'samples=samples.yaml',
         )
     )
+    # Need to have the sample template visible to the generator.
+    g._env.loader = jinja2.DictLoader({'sample.py.j2': ''})
+
     api_schema = make_api(naming=naming.Naming(name='Mollusc', version='v6'))
     actual_response = g.get_response(api_schema)
     expected_response = CodeGeneratorResponse(
@@ -393,6 +421,9 @@ def test_samplegen_id_disambiguation(mock_gmtime, mock_generate_sample, fs):
         )
     )
     g = generator.Generator(options.Options.build('samples=samples.yaml'))
+    # Need to have the sample template visible to the generator.
+    g._env.loader = jinja2.DictLoader({'sample.py.j2': ''})
+
     api_schema = make_api(naming=naming.Naming(name='Mollusc', version='v6'))
     actual_response = g.get_response(api_schema)
     expected_response = CodeGeneratorResponse(