Add Aspects to Bazel py_proto_library and py_grpc_library Rules (#27275)

* Move most py_proto_library macro contents into rule.

The _generate_pb2_src rule is modified to produce a valid PyInfo
provider that can be depended on directly by py_.* rules.

Keeping an intermediate macro in place increases complexity without
adding any real value.

* Move most py_grpc_library macro contents into rule.

The _generate_pb2_grpc_src rule is modified to produce a valid PyInfo
provider that can be depended on directly by py_.* rules.

Keeping an intermediate macro in place increase complexity without
adding any real value.

* Remove support for additional plugins.

It's easy enough for a plugin author to build their own rule-set, and it
adds extra unused (and untested) complexity to gRPC's rules.

* Rewrite py_proto_library to track dependencies correctly.

* Propagate Python dependencies of py_grpc_library rules to consumers.

* Add test coverage for aspect-based py_proto_library.

bazel/protobuf.bzl

@@ -34,6 +34,9 @@ def well_known_proto_libs():
         "@com_google_protobuf//:wrappers_proto",
     ]
 
+def is_well_known(label):
+    return label in well_known_proto_libs()
+
 def get_proto_root(workspace_root):
     """Gets the root protobuf directory.
 
@@ -141,8 +144,18 @@ def get_plugin_args(
         ),
     ]
 
-def _get_staged_proto_file(context, source_file):
-    if source_file.dirname == context.label.package or \
+def get_staged_proto_file(label, context, source_file):
+    """Copies a proto file to the appropriate location if necessary.
+
+    Args:
+      label: The label of the rule using the .proto file.
+      context: The ctx object for the rule or aspect.
+      source_file: The original .proto file.
+
+    Returns:
+      The original proto file OR a new file in the staged location.
+    """
+    if source_file.dirname == label.package or \
        is_in_virtual_imports(source_file):
         # Current target and source_file are in same package
         return source_file
@@ -170,7 +183,7 @@ def protos_from_context(context):
     protos = []
     for src in context.attr.deps:
         for file in src[ProtoInfo].direct_sources:
-            protos.append(_get_staged_proto_file(context, file))
+            protos.append(get_staged_proto_file(context.label, context, file))
     return protos
 
 def includes_from_deps(deps):

bazel/python_rules.bzl

@@ -21,42 +21,62 @@ load(
     "get_out_dir",
     "get_plugin_args",
     "get_proto_arguments",
+    "get_staged_proto_file",
     "includes_from_deps",
+    "is_in_virtual_imports",
+    "is_well_known",
     "protos_from_context",
 )
 
 _GENERATED_PROTO_FORMAT = "{}_pb2.py"
 _GENERATED_GRPC_PROTO_FORMAT = "{}_pb2_grpc.py"
 
-def _generate_py_impl(context):
-    protos = protos_from_context(context)
-    includes = includes_from_deps(context.attr.deps)
+PyProtoInfo = provider(
+    "The Python outputs from the Protobuf compiler.",
+    fields = {
+        "py_info": "A PyInfo provider for the generated code.",
+        "generated_py_srcs": "The direct (not transitive) generated Python source files.",
+    },
+)
+
+def _merge_pyinfos(pyinfos):
+    return PyInfo(
+        transitive_sources = depset(transitive = [p.transitive_sources for p in pyinfos]),
+        imports = depset(transitive = [p.imports for p in pyinfos]),
+    )
+
+def _gen_py_aspect_impl(target, context):
+    # Early return for well-known protos.
+    if is_well_known(str(context.label)):
+        return [
+            PyProtoInfo(py_info = context.attr._protobuf_library[PyInfo]),
+        ]
+
+    protos = []
+    for p in target[ProtoInfo].direct_sources:
+        protos.append(get_staged_proto_file(target.label, context, p))
+
+    includes = depset(direct = protos, transitive = [target[ProtoInfo].transitive_imports])
     out_files = declare_out_files(protos, context, _GENERATED_PROTO_FORMAT)
+    generated_py_srcs = out_files
+
     tools = [context.executable._protoc]
 
     out_dir = get_out_dir(protos, context)
+
     arguments = ([
         "--python_out={}".format(out_dir.path),
     ] + [
         "--proto_path={}".format(get_include_directory(i))
-        for i in includes
+        for i in includes.to_list()
     ] + [
         "--proto_path={}".format(context.genfiles_dir.path),
     ])
-    if context.attr.plugin:
-        arguments += get_plugin_args(
-            context.executable.plugin,
-            [],
-            out_dir.path,
-            False,
-            context.attr.plugin.label.name,
-        )
-        tools.append(context.executable.plugin)
 
     arguments += get_proto_arguments(protos, context.genfiles_dir.path)
 
     context.actions.run(
-        inputs = protos + includes,
+        inputs = protos + includes.to_list(),
         tools = tools,
         outputs = out_files,
         executable = context.executable._protoc,
@@ -68,26 +88,77 @@ def _generate_py_impl(context):
     if out_dir.import_path:
         imports.append("%s/%s/%s" % (context.workspace_name, context.label.package, out_dir.import_path))
 
+    py_info = PyInfo(transitive_sources = depset(direct = out_files), imports = depset(direct = imports))
+    return PyProtoInfo(
+        py_info = _merge_pyinfos(
+            [
+                py_info,
+                context.attr._protobuf_library[PyInfo],
+            ] + [dep[PyProtoInfo].py_info for dep in context.rule.attr.deps],
+        ),
+        generated_py_srcs = generated_py_srcs,
+    )
+
+_gen_py_aspect = aspect(
+    implementation = _gen_py_aspect_impl,
+    attr_aspects = ["deps"],
+    fragments = ["py"],
+    attrs = {
+        "_protoc": attr.label(
+            default = Label("//external:protocol_compiler"),
+            providers = ["files_to_run"],
+            executable = True,
+            cfg = "host",
+        ),
+        "_protobuf_library": attr.label(
+            default = Label("@com_google_protobuf//:protobuf_python"),
+            providers = [PyInfo],
+        ),
+    },
+)
+
+def _generate_py_impl(context):
+    if (len(context.attr.deps) != 1):
+        fail("Can only compile a single proto at a time.")
+
+    py_sources = []
+
+    # If the proto_library this rule *directly* depends on is in another
+    # package, then we generate .py files to import them in this package. This
+    # behavior is needed to allow rearranging of import paths to make Bazel
+    # outputs align with native python workflows.
+    #
+    # Note that this approach is vulnerable to protoc defining __all__ or other
+    # symbols with __ prefixes that need to be directly imported. Since these
+    # names are likely to be reserved for private APIs, the risk is minimal.
+    if context.label.package != context.attr.deps[0].label.package:
+        for py_src in context.attr.deps[0][PyProtoInfo].generated_py_srcs:
+            reimport_py_file = context.actions.declare_file(py_src.basename)
+            py_sources.append(reimport_py_file)
+            import_line = "from %s import *" % py_src.short_path.replace("/", ".")[:-len(".py")]
+            context.actions.write(reimport_py_file, import_line)
+
+    # Collect output PyInfo provider.
+    imports = [context.label.package + "/" + i for i in context.attr.imports]
+    py_info = PyInfo(transitive_sources = depset(direct = py_sources), imports = depset(direct = imports))
+    out_pyinfo = _merge_pyinfos([py_info, context.attr.deps[0][PyProtoInfo].py_info])
+
+    runfiles = context.runfiles(files = out_pyinfo.transitive_sources.to_list()).merge(context.attr._protobuf_library[DefaultInfo].data_runfiles)
     return [
-        DefaultInfo(files = depset(direct = out_files)),
-        PyInfo(
-            transitive_sources = depset(),
-            imports = depset(direct = imports),
+        DefaultInfo(
+            files = out_pyinfo.transitive_sources,
+            runfiles = runfiles,
         ),
+        out_pyinfo,
     ]
 
-_generate_pb2_src = rule(
+py_proto_library = rule(
     attrs = {
         "deps": attr.label_list(
             mandatory = True,
             allow_empty = False,
             providers = [ProtoInfo],
-        ),
-        "plugin": attr.label(
-            mandatory = False,
-            executable = True,
-            providers = ["files_to_run"],
-            cfg = "host",
+            aspects = [_gen_py_aspect],
         ),
         "_protoc": attr.label(
             default = Label("//external:protocol_compiler"),
@@ -95,46 +166,15 @@ _generate_pb2_src = rule(
             executable = True,
             cfg = "host",
         ),
+        "_protobuf_library": attr.label(
+            default = Label("@com_google_protobuf//:protobuf_python"),
+            providers = [PyInfo],
+        ),
+        "imports": attr.string_list(),
     },
     implementation = _generate_py_impl,
 )
 
-def py_proto_library(
-        name,
-        deps,
-        plugin = None,
-        **kwargs):
-    """Generate python code for a protobuf.
-
-    Args:
-      name: The name of the target.
-      deps: A list of proto_library dependencies. Must contain a single element.
-      plugin: An optional custom protoc plugin to execute together with
-        generating the protobuf code.
-      **kwargs: Additional arguments to be supplied to the invocation of
-        py_library.
-    """
-    codegen_target = "_{}_codegen".format(name)
-    if len(deps) != 1:
-        fail("Can only compile a single proto at a time.")
-
-    _generate_pb2_src(
-        name = codegen_target,
-        deps = deps,
-        plugin = plugin,
-        **kwargs
-    )
-
-    native.py_library(
-        name = name,
-        srcs = [":{}".format(codegen_target)],
-        deps = [
-            "@com_google_protobuf//:protobuf_python",
-            ":{}".format(codegen_target),
-        ],
-        **kwargs
-    )
-
 def _generate_pb2_grpc_src_impl(context):
     protos = protos_from_context(context)
     includes = includes_from_deps(context.attr.deps)
@@ -151,15 +191,6 @@ def _generate_pb2_grpc_src_impl(context):
         out_dir.path,
         False,
     )
-    if context.attr.plugin:
-        arguments += get_plugin_args(
-            context.executable.plugin,
-            [],
-            out_dir.path,
-            False,
-            context.attr.plugin.label.name,
-        )
-        tools.append(context.executable.plugin)
 
     arguments += [
         "--proto_path={}".format(get_include_directory(i))
@@ -177,13 +208,22 @@ def _generate_pb2_grpc_src_impl(context):
         mnemonic = "ProtocInvocation",
     )
 
+    p = PyInfo(transitive_sources = depset(direct = out_files))
+    py_info = _merge_pyinfos(
+        [
+            p,
+            context.attr._grpc_library[PyInfo],
+        ] + [dep[PyInfo] for dep in context.attr.py_deps],
+    )
+
+    runfiles = context.runfiles(files = out_files, transitive_files = py_info.transitive_sources).merge(context.attr._grpc_library[DefaultInfo].data_runfiles)
+
     return [
-        DefaultInfo(files = depset(direct = out_files)),
-        PyInfo(
-            transitive_sources = depset(),
-            # Imports are already configured by the generated py impl
-            imports = depset(),
+        DefaultInfo(
+            files = depset(direct = out_files),
+            runfiles = runfiles,
         ),
+        py_info,
     ]
 
 _generate_pb2_grpc_src = rule(
@@ -193,13 +233,12 @@ _generate_pb2_grpc_src = rule(
             allow_empty = False,
             providers = [ProtoInfo],
         ),
-        "strip_prefixes": attr.string_list(),
-        "plugin": attr.label(
-            mandatory = False,
-            executable = True,
-            providers = ["files_to_run"],
-            cfg = "host",
+        "py_deps": attr.label_list(
+            mandatory = True,
+            allow_empty = False,
+            providers = [PyInfo],
         ),
+        "strip_prefixes": attr.string_list(),
         "_grpc_plugin": attr.label(
             executable = True,
             providers = ["files_to_run"],
@@ -212,6 +251,10 @@ _generate_pb2_grpc_src = rule(
             cfg = "host",
             default = Label("//external:protocol_compiler"),
         ),
+        "_grpc_library": attr.label(
+            default = Label("//src/python/grpcio/grpc:grpcio"),
+            providers = [PyInfo],
+        ),
     },
     implementation = _generate_pb2_grpc_src_impl,
 )
@@ -220,7 +263,6 @@ def py_grpc_library(
         name,
         srcs,
         deps,
-        plugin = None,
         strip_prefixes = [],
         **kwargs):
     """Generate python code for gRPC services defined in a protobuf.
@@ -235,12 +277,9 @@ def py_grpc_library(
         stripped from the beginning of foo_pb2 modules imported by the
         generated stubs. This is useful in combination with the `imports`
         attribute of the `py_library` rule.
-      plugin: An optional custom protoc plugin to execute together with
-        generating the gRPC code.
       **kwargs: Additional arguments to be supplied to the invocation of
         py_library.
     """
-    codegen_grpc_target = "_{}_grpc_codegen".format(name)
     if len(srcs) != 1:
         fail("Can only compile a single proto at a time.")
 
@@ -248,23 +287,10 @@ def py_grpc_library(
         fail("Deps must have length 1.")
 
     _generate_pb2_grpc_src(
-        name = codegen_grpc_target,
+        name = name,
         deps = srcs,
+        py_deps = deps,
         strip_prefixes = strip_prefixes,
-        plugin = plugin,
-        **kwargs
-    )
-
-    native.py_library(
-        name = name,
-        srcs = [
-            ":{}".format(codegen_grpc_target),
-        ],
-        deps = [
-            Label("//src/python/grpcio/grpc:grpcio"),
-        ] + deps + [
-            ":{}".format(codegen_grpc_target),
-        ],
         **kwargs
     )
 

bazel/test/python_test_repo/BUILD

@@ -21,7 +21,9 @@ load(
     "py_proto_library",
 )
 
-package(default_testonly = 1)
+package(
+    default_testonly = 1,
+)
 
 proto_library(
     name = "helloworld_proto",
@@ -82,10 +84,8 @@ proto_library(
     ],
 )
 
-# Also test the custom plugin execution parameter
 py_proto_library(
     name = "helloworld_moved_py_pb2",
-    plugin = ":dummy_plugin",
     deps = [":helloworld_moved_proto"],
 )
 
@@ -108,11 +108,42 @@ py_test(
     ],
 )
 
-py_binary(
-    name = "dummy_plugin",
-    srcs = [":dummy_plugin.py"],
+# Test that a py_proto_library wrapping a proto_library in another package can
+# be imported from the package that contains the py_proto_library *AND* from
+# the package that contains the proto_library.
+py_proto_library(
+    name = "subpackage_py_pb2",
+    deps = ["//in_subpackage:subpackage_proto"],
+)
+
+py_test(
+    name = "import_from_this_package_subpackage_test",
+    srcs = ["import_from_this_package.py"],
+    main = "import_from_this_package.py",
+    python_version = "PY3",
+    deps = [
+        ":subpackage_py_pb2",
+    ],
+)
+
+py_test(
+    name = "import_from_proto_library_package_test",
+    srcs = ["import_from_proto_library_package.py"],
+    main = "import_from_proto_library_package.py",
     python_version = "PY3",
     deps = [
-        "@com_google_protobuf//:protobuf_python",
+        ":subpackage_py_pb2",
+    ],
+)
+
+# Test that a py_proto_library can be successfully imported without requiring
+# explicit dependencies on unused dependencies of the generated code.
+py_test(
+    name = "transitive_proto_dep_test",
+    srcs = ["transitive_proto_dep.py"],
+    main = "transitive_proto_dep.py",
+    python_version = "PY3",
+    deps = [
+        ":helloworld_py_pb2",
     ],
 )

bazel/test/python_test_repo/dummy_plugin.py

@@ -1,37 +0,0 @@
-# Copyright 2019 the gRPC authors.
-#
-# 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
-#
-#     http://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.
-"""A dummy plugin for testing"""
-
-import sys
-
-from google.protobuf.compiler.plugin_pb2 import CodeGeneratorRequest
-from google.protobuf.compiler.plugin_pb2 import CodeGeneratorResponse
-
-
-def main(input_file=sys.stdin, output_file=sys.stdout):
-    request = CodeGeneratorRequest.FromString(input_file.buffer.read())
-    answer = []
-    for fname in request.file_to_generate:
-        answer.append(CodeGeneratorResponse.File(
-            name=fname.replace('.proto', '_pb2.py'),
-            insertion_point='module_scope',
-            content="# Hello {}, I'm a dummy plugin!".format(fname),
-        ))
-
-    cgr = CodeGeneratorResponse(file=answer)
-    output_file.buffer.write(cgr.SerializeToString())
-
-
-if __name__ == '__main__':
-    main()

bazel/test/python_test_repo/import_from_proto_library_package.py

@@ -0,0 +1,18 @@
+# Copyright 2021 the gRPC authors.
+#
+# 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
+#
+#     http://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.
+"""A trivial executable that imports Protobuf generated code where the
+proto_library and py_proto_library are in different Bazel packages.
+"""
+
+import in_subpackage.subpackage_pb2

bazel/test/python_test_repo/import_from_this_package.py

@@ -0,0 +1,18 @@
+# Copyright 2021 the gRPC authors.
+#
+# 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
+#
+#     http://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.
+"""A trivial executable that imports Protobuf generated code where the
+proto_library and py_proto_library are in different Bazel packages.
+"""
+
+import subpackage_pb2

bazel/test/python_test_repo/in_subpackage/BUILD

@@ -0,0 +1,27 @@
+# gRPC Bazel BUILD file.
+#
+# Copyright 2021 The gRPC authors.
+#
+# 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
+#
+#     http://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.
+
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+package(
+    default_testonly = 1,
+    default_visibility = ["//:__subpackages__"],
+)
+
+proto_library(
+    name = "subpackage_proto",
+    srcs = ["subpackage.proto"],
+)

bazel/test/python_test_repo/in_subpackage/subpackage.proto

@@ -0,0 +1,27 @@
+// Copyright 2021 The gRPC authors.
+//
+// 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
+//
+//     http://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.
+
+syntax = "proto3";
+
+option java_multiple_files = true;
+option java_package = "io.grpc.examples.subpackage";
+option java_outer_classname = "SubpackageProto";
+option objc_class_prefix = "SP";
+
+package subpackage;
+
+message Subpackaged {
+  string name = 1;
+}
+

bazel/test/python_test_repo/transitive_proto_dep.py

@@ -0,0 +1,17 @@
+# Copyright 2021 the gRPC authors.
+#
+# 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
+#
+#     http://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.
+"""A trivial executable that imports helloworld_pb2, but NOT its transitive
+dependencies."""
+
+import helloworld_pb2

src/python/grpcio_channelz/grpc_channelz/v1/BUILD.bazel

@@ -17,6 +17,7 @@ package(default_visibility = ["//visibility:public"])
 
 py_proto_library(
     name = "channelz_py_pb2",
+    imports = ["../../"],
     deps = ["//src/proto/grpc/channelz:channelz_proto_descriptors"],
 )
 

src/python/grpcio_health_checking/grpc_health/v1/BUILD.bazel

@@ -17,6 +17,7 @@ package(default_visibility = ["//visibility:public"])
 
 py_proto_library(
     name = "health_py_pb2",
+    imports = ["../../"],
     deps = ["//src/proto/grpc/health/v1:health_proto_descriptor"],
 )
 

src/python/grpcio_reflection/grpc_reflection/v1alpha/BUILD.bazel

@@ -18,6 +18,7 @@ package(default_visibility = ["//visibility:public"])
 
 py_proto_library(
     name = "reflection_py_pb2",
+    imports = ["../../"],
     deps = ["//src/proto/grpc/reflection/v1alpha:reflection_proto_descriptor"],
 )