feat: add generator option proto-plus-dep (#1596)

* feat: add generator option proto-plus-dep

* fix mypy

* fix style

* update goldens

* fix tests

* clean up

* clean up

* add support for dependency google-cloud-kms

* Address review comments

* add missing file

* run formatting tool

* style

* address review comments

* formatting

* address review feedback

* style

* add pypi package google-geo-type

* bump google-cloud-documentai dependency to 2.0.0

* address review comment

* add docstring

* address review comment

* fix docs

Author: parthea

packages/gapic-generator/gapic/ads-templates/%namespace/%name/%version/%sub/services/%service/transports/rest.py.j2

@@ -322,10 +322,10 @@ class {{service.name}}RestTransport({{service.name}}Transport):
             {% endfor %}{# rule in method.http_options #}
             ]
             request, metadata = self._interceptor.pre_{{ method.name|snake_case }}(request, metadata)
-            {% if method.input.ident.is_external_type %}
-            pb_request = request
-            {% else %}
+            {% if method.input.ident.is_proto_plus_type %}
             pb_request = {{method.input.ident}}.pb(request)
+            {% else %}
+            pb_request = request
             {% endif %}
             transcoded_request = path_template.transcode(http_options, pb_request)
 
@@ -384,10 +384,10 @@ class {{service.name}}RestTransport({{service.name}}Transport):
             resp = rest_streaming.ResponseIterator(response, {{method.output.ident}})
             {% else %}
             resp = {{method.output.ident}}()
-            {% if method.output.ident.is_external_type %}
-            pb_resp = resp
-            {% else %}
+            {% if method.output.ident.is_proto_plus_type %}
             pb_resp = {{method.output.ident}}.pb(resp)
+            {% else %}
+            pb_resp = resp
             {% endif %}
 
             json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True)

packages/gapic-generator/gapic/ads-templates/tests/unit/gapic/%name_%version/%sub/test_%service.py.j2

@@ -1118,10 +1118,10 @@ def test_{{ method_name }}_rest_required_fields(request_type={{ method.input.ide
     {% endif %}{# default is str #}
     {% endfor %}
     request = request_type(**request_init)
-    {% if method.input.ident.is_external_type %}
-    pb_request = request
-    {% else %}
+    {% if method.input.ident.is_proto_plus_type %}
     pb_request = request_type.pb(request)
+    {% else %}
+    pb_request = request
     {% endif %}
     jsonified_request = json.loads(json_format.MessageToJson(
         pb_request,
@@ -1189,10 +1189,10 @@ def test_{{ method_name }}_rest_required_fields(request_type={{ method.input.ide
         with mock.patch.object(path_template, 'transcode') as transcode:
             # A uri without fields and an empty body will force all the
             # request fields to show up in the query_params.
-            {% if method.input.ident.is_external_type %}
-            pb_request = request
-            {% else %}
+            {% if method.input.ident.is_proto_plus_type %}
             pb_request = request_type.pb(request)
+            {% else %}
+            pb_request = request
             {% endif %}
             transcode_result = {
                 'uri': 'v1/sample_method',
@@ -1212,10 +1212,10 @@ def test_{{ method_name }}_rest_required_fields(request_type={{ method.input.ide
             json_return_value = json_format.MessageToJson(return_value)
             {% else %}
 
-            {% if method.output.ident.is_external_type %}
-            pb_return_value = return_value
-            {% else %}
+            {% if method.output.ident.is_proto_plus_type %}
             pb_return_value = {{ method.output.ident }}.pb(return_value)
+            {% else %}
+            pb_return_value = return_value
             {% endif %}
             json_return_value = json_format.MessageToJson(pb_return_value)
             {% endif %}
@@ -1287,10 +1287,10 @@ def test_{{ method_name }}_rest_interceptors(null_interceptor):
         {% if not method.void %}
         post.assert_not_called()
         {% endif %}
-        {% if method.input.ident.is_external_type %}
-        pb_message = {{ method.input.ident }}()
-        {% else %}
+        {% if method.input.ident.is_proto_plus_type %}
         pb_message = {{ method.input.ident }}.pb({{ method.input.ident }}())
+        {% else %}
+        pb_message = {{ method.input.ident }}()
         {% endif %}
         transcode.return_value = {
             "method": "post",

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

@@ -27,6 +27,7 @@
 """
 
 import dataclasses
+import re
 from typing import FrozenSet, Tuple, Optional
 
 from google.protobuf import descriptor_pb2
@@ -89,10 +90,8 @@ def __str__(self) -> str:
             if self.module_alias:
                 module_name = self.module_alias
 
-            # This module is from a different proto package
-            # Most commonly happens for a common proto
-            # https://pypi.org/project/googleapis-common-protos/
-            if self.is_external_type:
+            # Add _pb2 suffix except when it is a proto-plus type
+            if not self.is_proto_plus_type:
                 module_name = f'{self.module}_pb2'
 
             # Return the dot-separated Python identifier.
@@ -103,8 +102,23 @@ def __str__(self) -> str:
         return '.'.join(self.parent + (self.name,))
 
     @property
-    def is_external_type(self):
-        return not self.proto_package.startswith(self.api_naming.proto_package)
+    def is_proto_plus_type(self) -> bool:
+        """This function is used to determine whether a given package `self.proto_package`
+        is using proto-plus types or protobuf types. There are 2 scenarios where the package
+        is expected to use proto-plus types:
+        1) When `self.proto_package` starts with `self.api_naming.proto_package`, then
+        the given package has the same namespace as the one that is being generated. It is assumed
+        that the gapic generator always generates packages with proto-plus types.
+        2) When `self.proto_package` is explicitly in `self.api_naming.proto_plus_deps` which is
+        populated via the generator option `proto-plus-deps`.
+
+        Returns:
+            bool: Whether the given package uses proto-plus types or not.
+        """
+        return self.proto_package.startswith(self.api_naming.proto_package) or (
+            hasattr(self.api_naming, "proto_plus_deps")
+            and self.proto_package in self.api_naming.proto_plus_deps
+        )
 
     @cached_property
     def __cached_string_repr(self):
@@ -188,6 +202,29 @@ def python_import(self) -> imp.Import:
                 alias=self.module_alias,
             )
 
+        if self.is_proto_plus_type:
+            # We need to change the import statement to use an
+            # underscore between the module and the version. For example,
+            # change google.cloud.documentai.v1 to google.cloud.documentai_v1.
+            # Check if the package name contains a version.
+            version_regex = "^v\d[^/]*$"
+            regex_match = re.match(version_regex, self.package[-1])
+
+            if regex_match and len(self.package) > 1:
+                versioned_module = f"{self.package[-2]}_{regex_match[0]}"
+                return imp.Import(
+                    package=self.package[:-2] +
+                    (versioned_module, 'types'),
+                    module=self.module,
+                    alias=self.module_alias,
+                )
+            else:
+                return imp.Import(
+                    package=self.package + ('types',),
+                    module=self.module,
+                    alias=self.module_alias,
+                )
+
         # Return the standard import.
         return imp.Import(
             package=self.package,

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

@@ -43,6 +43,7 @@ class Naming(abc.ABC):
     product_name: str = ''
     proto_package: str = ''
     _warehouse_package_name: str = ''
+    proto_plus_deps: Tuple[str, ...] = dataclasses.field(default_factory=tuple)
 
     def __post_init__(self):
         if not self.product_name:
@@ -146,6 +147,11 @@ def build(
             package_info = dataclasses.replace(package_info,
                 _warehouse_package_name=opts.warehouse_package_name
                                                )
+        if opts.proto_plus_deps:
+            package_info = dataclasses.replace(
+                package_info,
+                proto_plus_deps=opts.proto_plus_deps,
+            )
 
         # Done; return the naming information.
         return package_info

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

@@ -77,7 +77,7 @@ def __hash__(self):
     def name(self) -> str:
         """Used to prevent collisions with python keywords"""
         name = self.field_pb.name
-        return name + "_" if name in utils.RESERVED_NAMES and not self.meta.address.is_external_type else name
+        return name + "_" if name in utils.RESERVED_NAMES and self.meta.address.is_proto_plus_type else name
 
     @utils.cached_property
     def ident(self) -> metadata.FieldIdentifier:

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

@@ -356,10 +356,10 @@ class {{service.name}}RestTransport({{service.name}}Transport):
             {% endfor %}{# rule in method.http_options #}
             ]
             request, metadata = self._interceptor.pre_{{ method.name|snake_case }}(request, metadata)
-            {% if method.input.ident.is_external_type %}
-            pb_request = request
-            {% else %}
+            {% if method.input.ident.is_proto_plus_type %}
             pb_request = {{method.input.ident}}.pb(request)
+            {% else %}
+            pb_request = request
             {% endif %}
             transcoded_request = path_template.transcode(http_options, pb_request)
 
@@ -418,10 +418,10 @@ class {{service.name}}RestTransport({{service.name}}Transport):
             resp = rest_streaming.ResponseIterator(response, {{method.output.ident}})
             {% else %}
             resp = {{method.output.ident}}()
-            {% if method.output.ident.is_external_type %}
-            pb_resp = resp
-            {% else %}
+            {% if method.output.ident.is_proto_plus_type %}
             pb_resp = {{method.output.ident}}.pb(resp)
+            {% else %}
+            pb_resp = resp
             {% endif %}
 
             json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True)

packages/gapic-generator/gapic/templates/_pypi_packages.j2

@@ -0,0 +1,15 @@
+<!--
+Dict of PyPI packages available by API in the form <tuple>, <tuple>
+Where the first <tuple> contains the API and the second tuple contains the
+PyPI package name, the minimum allowed version and the maximum allowed version.
+
+Note: Set the minimum version for google-cloud-documentai to 2.0.0 which has support for `barcode` in `google.cloud.documentai.types`
+-->
+{% set pypi_packages = {
+    ("google", "apps", "script", "type"): {"package_name": "google-apps-script-type", "lower_bound": "0.2.0", "upper_bound": "1.0.0dev"},
+    ("google", "geo", "type"): {"package_name": "google-geo-type", "lower_bound": "0.1.0", "upper_bound": "1.0.0dev"},
+    ("google", "cloud", "documentai", "v1"): {"package_name": "google-cloud-documentai", "lower_bound": "2.0.0", "upper_bound": "3.0.0dev"},
+    ("google", "cloud", "kms", "v1"): {"package_name": "google-cloud-kms", "lower_bound": "2.3.0", "upper_bound": "3.0.0dev"},
+    ("google", "iam", "v1"): {"package_name": "grpc-google-iam-v1", "lower_bound": "0.12.4", "upper_bound": "1.0.0dev"}
+}
+%}

packages/gapic-generator/gapic/templates/setup.py.j2

@@ -1,5 +1,5 @@
 {% extends '_base.py.j2' %}
-
+{% from '_pypi_packages.j2' import pypi_packages %}
 {% block content %}
 
 import io
@@ -31,13 +31,14 @@ dependencies = [
     "proto-plus >= 1.22.0, <2.0.0dev",
     "proto-plus >= 1.22.2, <2.0.0dev; python_version>='3.11'",
     "protobuf>=3.19.5,<5.0.0dev,!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5",
-    {# TODO: Remove after https://github.com/googleapis/gapic-generator-python/pull/1240 is merged. #}
-    {% if api.requires_package(('google', 'iam', 'v1')) or opts.add_iam_methods or api.has_iam_mixin %}
-        'grpc-google-iam-v1 >= 0.12.4, < 1.0.0dev',
+    {% for package_tuple, package_info in pypi_packages.items() %}
+    {# Quick check to make sure the package is different from this setup.py #}
+    {% if api.naming.warehouse_package_name != package_info.package_name %}
+    {% if api.requires_package(package_tuple) %}
+    "{{ package_info.package_name }} >= {{ package_info.lower_bound }}, <{{ package_info.upper_bound }}",
     {% endif %}
-    {% if api.requires_package(('google', 'cloud', 'documentai', 'v1')) %}
-        'google-cloud-documentai >= 1.2.1, < 3.0.0dev',
     {% endif %}
+    {% endfor %}
 ]
 url = "https://github.com/googleapis/python-{{ api.naming.warehouse_package_name|replace("google-cloud-", "")|replace("google-", "") }}"
 

packages/gapic-generator/gapic/templates/testing/_default_constraints.j2

@@ -1,11 +1,14 @@
+{% from '_pypi_packages.j2' import pypi_packages %}
 # This constraints file is required for unit tests.
 # List all library dependencies and extras in this file.
 google-api-core
 proto-plus
 protobuf
-{% if api.requires_package(('google', 'iam', 'v1')) or opts.add_iam_methods or api.has_iam_mixin %}
-grpc-google-iam-v1
+{% for package_tuple, package_info in pypi_packages.items() %}
+{# Quick check to make sure the package is different from this setup.py #}
+{% if api.naming.warehouse_package_name != package_info.package_name %}
+{% if api.requires_package(package_tuple) %}
+{{ package_info.package_name }}
 {% endif %}
-{% if api.requires_package(('google', 'cloud', 'documentai', 'v1')) %}
-google-cloud-documentai
-{% endif %}
+{% endif %}
+{% endfor %}

packages/gapic-generator/gapic/templates/testing/constraints-3.7.txt.j2

@@ -1,3 +1,4 @@
+{% from '_pypi_packages.j2' import pypi_packages %}
 # This constraints file is used to check that lower bounds
 # are correct in setup.py
 # List all library dependencies and extras in this file.
@@ -7,9 +8,11 @@
 google-api-core==1.34.0
 proto-plus==1.22.0
 protobuf==3.19.5
-{% if api.requires_package(('google', 'iam', 'v1')) or opts.add_iam_methods or api.has_iam_mixin %}
-grpc-google-iam-v1==0.12.4
+{% for package_tuple, package_info in pypi_packages.items() %}
+{# Quick check to make sure the package is different from this setup.py #}
+{% if api.naming.warehouse_package_name != package_info.package_name %}
+{% if api.requires_package(package_tuple) %}
+{{ package_info.package_name }}=={{ package_info.lower_bound }}
 {% endif %}
-{% if api.requires_package(('google', 'cloud', 'documentai', 'v1')) %}
-google-cloud-documentai==1.2.1
 {% endif %}
+{% endfor %}