Merge pull request #20846 from vam-google/master

[bazel] Add an ability to call an optional custom plugin for py_proto_library and py_grpc_library

bazel/protobuf.bzl

@@ -89,7 +89,12 @@ def get_include_directory(source_file):
     else:
         return source_file.root.path if source_file.root.path else "."
 
-def get_plugin_args(plugin, flags, dir_out, generate_mocks):
+def get_plugin_args(
+        plugin,
+        flags,
+        dir_out,
+        generate_mocks,
+        plugin_name = "PLUGIN"):
     """Returns arguments configuring protoc to use a plugin for a language.
 
     Args:
@@ -97,16 +102,28 @@ def get_plugin_args(plugin, flags, dir_out, generate_mocks):
       flags: The plugin flags to be passed to protoc.
       dir_out: The output directory for the plugin.
       generate_mocks: A bool indicating whether to generate mocks.
-
+      plugin_name: A name of the plugin, it is required to be unique when there
+      are more than one plugin used in a single protoc command.
     Returns:
       A list of protoc arguments configuring the plugin.
     """
     augmented_flags = list(flags)
     if generate_mocks:
         augmented_flags.append("generate_mock_code=true")
+
+    augmented_dir_out = dir_out
+    if augmented_flags:
+        augmented_dir_out = ",".join(augmented_flags) + ":" + dir_out
+
     return [
-        "--plugin=protoc-gen-PLUGIN=" + plugin.path,
-        "--PLUGIN_out=" + ",".join(augmented_flags) + ":" + dir_out,
+        "--plugin=protoc-gen-{plugin_name}={plugin_path}".format(
+            plugin_name = plugin_name,
+            plugin_path = plugin.path,
+        ),
+        "--{plugin_name}_out={dir_out}".format(
+            plugin_name = plugin_name,
+            dir_out = augmented_dir_out,
+        )
     ]
 
 def _get_staged_proto_file(context, source_file):

bazel/python_rules.bzl

@@ -29,6 +29,16 @@ def _generate_py_impl(context):
     ] + [
         "--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(
@@ -59,6 +69,12 @@ _generate_pb2_src = rule(
             allow_empty = False,
             providers = [ProtoInfo],
         ),
+        "plugin": attr.label(
+            mandatory = False,
+            executable = True,
+            providers = ["files_to_run"],
+            cfg = "host",
+        ),
         "_protoc": attr.label(
             default = Label("//external:protocol_compiler"),
             providers = ["files_to_run"],
@@ -72,12 +88,17 @@ _generate_pb2_src = rule(
 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:
@@ -87,6 +108,7 @@ def py_proto_library(
     _generate_pb2_src(
         name = codegen_target,
         deps = deps,
+        plugin = plugin,
         **kwargs
     )
 
@@ -108,14 +130,23 @@ def _generate_pb2_grpc_src_impl(context):
     plugin_flags = ["grpc_2_0"] + context.attr.strip_prefixes
 
     arguments = []
-    tools = [context.executable._protoc, context.executable._plugin]
+    tools = [context.executable._protoc, context.executable._grpc_plugin]
     out_dir = get_out_dir(protos, context)
     arguments += get_plugin_args(
-        context.executable._plugin,
+        context.executable._grpc_plugin,
         plugin_flags,
         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))
@@ -153,7 +184,13 @@ _generate_pb2_grpc_src = rule(
             providers = [ProtoInfo],
         ),
         "strip_prefixes": attr.string_list(),
-        "_plugin": attr.label(
+        "plugin": attr.label(
+            mandatory = False,
+            executable = True,
+            providers = ["files_to_run"],
+            cfg = "host",
+        ),
+        "_grpc_plugin": attr.label(
             executable = True,
             providers = ["files_to_run"],
             cfg = "host",
@@ -173,6 +210,7 @@ def py_grpc_library(
     name,
     srcs,
     deps,
+    plugin = None,
     strip_prefixes = [],
     **kwargs):
     """Generate python code for gRPC services defined in a protobuf.
@@ -187,6 +225,8 @@ 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.
     """
@@ -201,6 +241,7 @@ def py_grpc_library(
         name = codegen_grpc_target,
         deps = srcs,
         strip_prefixes = strip_prefixes,
+        plugin = plugin,
         **kwargs
     )
 

bazel/test/python_test_repo/BUILD

@@ -79,9 +79,11 @@ proto_library(
     strip_import_prefix = ""
 )
 
+# Also test the custom plugin execution parameter
 py_proto_library(
     name = "helloworld_moved_py_pb2",
     deps = [":helloworld_moved_proto"],
+    plugin = ":dummy_plugin"
 )
 
 py_grpc_library(
@@ -101,3 +103,11 @@ py2and3_test(
         ":timestamp_py_pb2",
     ],
 )
+
+py_binary(
+    name = "dummy_plugin",
+    srcs = [":dummy_plugin.py"],
+    deps = [
+        "@com_google_protobuf//:protobuf_python",
+    ],
+)

bazel/test/python_test_repo/dummy_plugin.py

@@ -0,0 +1,37 @@
+# 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()