Merge pull request #20141 from gnossen/python_strip_prefixes

Add strip_prefix to python protoc plugin and py_grpc_library

bazel/python_rules.bzl

@@ -93,11 +93,13 @@ def _generate_pb2_grpc_src_impl(context):
     proto_root = get_proto_root(context.label.workspace_root)
     out_files = declare_out_files(protos, context, _GENERATED_GRPC_PROTO_FORMAT)
 
+    plugin_flags = ["grpc_2_0"] + context.attr.strip_prefixes
+
     arguments = []
     tools = [context.executable._protoc, context.executable._plugin]
     arguments += get_plugin_args(
         context.executable._plugin,
-        [],
+        plugin_flags,
         context.genfiles_dir.path,
         False,
     )
@@ -127,6 +129,7 @@ _generate_pb2_grpc_src = rule(
             allow_empty = False,
             providers = [ProtoInfo],
         ),
+        "strip_prefixes": attr.string_list(),
         "_plugin": attr.label(
             executable = True,
             providers = ["files_to_run"],
@@ -147,6 +150,7 @@ def py_grpc_library(
     name,
     srcs,
     deps,
+    strip_prefixes = [],
     **kwargs):
     """Generate python code for gRPC services defined in a protobuf.
 
@@ -156,6 +160,10 @@ def py_grpc_library(
         schema of the service.
       deps: (List of `labels`) a single py_proto_library target for the
         proto_library in `srcs`.
+      strip_prefixes: (List of `strings`) If provided, this prefix will be
+        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.
     """
     codegen_grpc_target = "_{}_grpc_codegen".format(name)
     if len(srcs) != 1:
@@ -167,6 +175,7 @@ def py_grpc_library(
     _generate_pb2_grpc_src(
         name = codegen_grpc_target,
         deps = srcs,
+        strip_prefixes = strip_prefixes,
         **kwargs
     )
 

examples/python/debug/get_stats.py

@@ -21,13 +21,8 @@ import logging
 import argparse
 import grpc
 
-# TODO(https://github.com/grpc/grpc/issues/19863): Remove.
-try:
-    from src.python.grpcio_channelz.grpc_channelz.v1 import channelz_pb2
-    from src.python.grpcio_channelz.grpc_channelz.v1 import channelz_pb2_grpc
-except ImportError:
-    from grpc_channelz.v1 import channelz_pb2
-    from grpc_channelz.v1 import channelz_pb2_grpc
+from grpc_channelz.v1 import channelz_pb2
+from grpc_channelz.v1 import channelz_pb2_grpc
 
 
 def run(addr):

src/compiler/python_generator.cc

@@ -756,6 +756,29 @@ static bool GenerateGrpc(GeneratorContext* context, PrivateGenerator& generator,
   }
 }
 
+static bool ParseParameters(const grpc::string& parameter,
+                            grpc::string* grpc_version,
+                            std::vector<grpc::string>* strip_prefixes,
+                            grpc::string* error) {
+  std::vector<grpc::string> comma_delimited_parameters;
+  grpc_python_generator::Split(parameter, ',', &comma_delimited_parameters);
+  if (comma_delimited_parameters.size() == 1 &&
+      comma_delimited_parameters[0].empty()) {
+    *grpc_version = "grpc_2_0";
+  } else if (comma_delimited_parameters.size() == 1) {
+    *grpc_version = comma_delimited_parameters[0];
+  } else if (comma_delimited_parameters.size() == 2) {
+    *grpc_version = comma_delimited_parameters[0];
+    std::copy(comma_delimited_parameters.begin() + 1,
+              comma_delimited_parameters.end(),
+              std::back_inserter(*strip_prefixes));
+  } else {
+    *error = "--grpc_python_out received too many comma-delimited parameters.";
+    return false;
+  }
+  return true;
+}
+
 bool PythonGrpcGenerator::Generate(const FileDescriptor* file,
                                    const grpc::string& parameter,
                                    GeneratorContext* context,
@@ -778,14 +801,19 @@ bool PythonGrpcGenerator::Generate(const FileDescriptor* file,
   generator_file_name = file->name();
 
   ProtoBufFile pbfile(file);
-  PrivateGenerator generator(config_, &pbfile);
-  if (parameter == "" || parameter == "grpc_2_0") {
+  grpc::string grpc_version;
+  GeneratorConfiguration extended_config(config_);
+  bool success = ParseParameters(parameter, &grpc_version,
+                                 &(extended_config.prefixes_to_filter), error);
+  PrivateGenerator generator(extended_config, &pbfile);
+  if (!success) return false;
+  if (grpc_version == "grpc_2_0") {
     return GenerateGrpc(context, generator, pb2_grpc_file_name, true);
-  } else if (parameter == "grpc_1_0") {
+  } else if (grpc_version == "grpc_1_0") {
     return GenerateGrpc(context, generator, pb2_grpc_file_name, true) &&
            GenerateGrpc(context, generator, pb2_file_name, false);
   } else {
-    *error = "Invalid parameter '" + parameter + "'.";
+    *error = "Invalid grpc version '" + grpc_version + "'.";
     return false;
   }
 }

src/compiler/python_generator_helpers.h

@@ -136,6 +136,16 @@ StringVector get_all_comments(const DescriptorType* descriptor) {
   return comments;
 }
 
+inline void Split(const grpc::string& s, char delim,
+                  std::vector<grpc::string>* append_to) {
+  auto current = s.begin();
+  while (current <= s.end()) {
+    auto next = std::find(current, s.end(), delim);
+    append_to->emplace_back(current, next);
+    current = next + 1;
+  }
+}
+
 }  // namespace
 
 }  // namespace grpc_python_generator

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

@@ -10,6 +10,7 @@ py_grpc_library(
     name = "channelz_py_pb2_grpc",
     srcs = ["//src/proto/grpc/channelz:channelz_proto_descriptors"],
     deps = [":channelz_py_pb2"],
+    strip_prefixes = ["src.python.grpcio_channelz."],
 )
 
 py_library(

src/python/grpcio_channelz/grpc_channelz/v1/channelz.py

@@ -16,13 +16,8 @@
 import grpc
 from grpc._cython import cygrpc
 
-# TODO(https://github.com/grpc/grpc/issues/19863): Remove.
-try:
-    from src.python.grpcio_channelz.grpc_channelz.v1 import channelz_pb2 as _channelz_pb2
-    from src.python.grpcio_channelz.grpc_channelz.v1 import channelz_pb2_grpc as _channelz_pb2_grpc
-except ImportError:
-    import grpc_channelz.v1.channelz_pb2 as _channelz_pb2
-    import grpc_channelz.v1.channelz_pb2_grpc as _channelz_pb2_grpc
+import grpc_channelz.v1.channelz_pb2 as _channelz_pb2
+import grpc_channelz.v1.channelz_pb2_grpc as _channelz_pb2_grpc
 
 from google.protobuf import json_format
 

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

@@ -10,6 +10,7 @@ py_grpc_library(
     name = "health_py_pb2_grpc",
     srcs = ["//src/proto/grpc/health/v1:health_proto_descriptor",],
     deps = [":health_py_pb2"],
+    strip_prefixes = ["src.python.grpcio_health_checking."],
 )
 
 py_library(

src/python/grpcio_health_checking/grpc_health/v1/health.py

@@ -18,13 +18,8 @@ import threading
 
 import grpc
 
-# TODO(https://github.com/grpc/grpc/issues/19863): Remove.
-try:
-    from src.python.grpcio_health_checking.grpc_health.v1 import health_pb2 as _health_pb2
-    from src.python.grpcio_health_checking.grpc_health.v1 import health_pb2_grpc as _health_pb2_grpc
-except ImportError:
-    from grpc_health.v1 import health_pb2 as _health_pb2
-    from grpc_health.v1 import health_pb2_grpc as _health_pb2_grpc
+from grpc_health.v1 import health_pb2 as _health_pb2
+from grpc_health.v1 import health_pb2_grpc as _health_pb2_grpc
 
 SERVICE_NAME = _health_pb2.DESCRIPTOR.services_by_name['Health'].full_name
 

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

@@ -12,6 +12,7 @@ py_grpc_library(
     name = "reflection_py_pb2_grpc",
     srcs = ["//src/proto/grpc/reflection/v1alpha:reflection_proto_descriptor",],
     deps = ["reflection_py_pb2"],
+    strip_prefixes = ["src.python.grpcio_reflection."],
 )
 
 py_library(

src/python/grpcio_reflection/grpc_reflection/v1alpha/reflection.py

@@ -17,15 +17,8 @@ import grpc
 from google.protobuf import descriptor_pb2
 from google.protobuf import descriptor_pool
 
-# TODO(https://github.com/grpc/grpc/issues/19863): Remove.
-try:
-    from src.python.grpcio_reflection.grpc_reflection.v1alpha \
-        import reflection_pb2 as _reflection_pb2
-    from src.python.grpcio_reflection.grpc_reflection.v1alpha \
-        import reflection_pb2_grpc as _reflection_pb2_grpc
-except ImportError:
-    from grpc_reflection.v1alpha import reflection_pb2 as _reflection_pb2
-    from grpc_reflection.v1alpha import reflection_pb2_grpc as _reflection_pb2_grpc
+from grpc_reflection.v1alpha import reflection_pb2 as _reflection_pb2
+from grpc_reflection.v1alpha import reflection_pb2_grpc as _reflection_pb2_grpc
 
 _POOL = descriptor_pool.Default()
 SERVICE_NAME = _reflection_pb2.DESCRIPTOR.services_by_name[

src/python/grpcio_tests/tests/channelz/_channelz_servicer_test.py

@@ -19,15 +19,9 @@ from concurrent import futures
 
 import grpc
 
-# TODO(https://github.com/grpc/grpc/issues/19863): Remove.
-try:
-    from src.python.grpcio_channelz.grpc_channelz.v1 import channelz
-    from src.python.grpcio_channelz.grpc_channelz.v1 import channelz_pb2
-    from src.python.grpcio_channelz.grpc_channelz.v1 import channelz_pb2_grpc
-except ImportError:
-    from grpc_channelz.v1 import channelz
-    from grpc_channelz.v1 import channelz_pb2
-    from grpc_channelz.v1 import channelz_pb2_grpc
+from grpc_channelz.v1 import channelz
+from grpc_channelz.v1 import channelz_pb2
+from grpc_channelz.v1 import channelz_pb2_grpc
 
 from tests.unit import test_common
 from tests.unit.framework.common import test_constants

src/python/grpcio_tests/tests/health_check/_health_servicer_test.py

@@ -20,15 +20,9 @@ import unittest
 
 import grpc
 
-# TODO(https://github.com/grpc/grpc/issues/19863): Remove.
-try:
-    from src.python.grpcio_health_checking.grpc_health.v1 import health
-    from src.python.grpcio_health_checking.grpc_health.v1 import health_pb2
-    from src.python.grpcio_health_checking.grpc_health.v1 import health_pb2_grpc
-except ImportError:
-    from grpc_health.v1 import health
-    from grpc_health.v1 import health_pb2
-    from grpc_health.v1 import health_pb2_grpc
+from grpc_health.v1 import health
+from grpc_health.v1 import health_pb2
+from grpc_health.v1 import health_pb2_grpc
 
 from tests.unit import test_common
 from tests.unit import thread_pool

src/python/grpcio_tests/tests/reflection/_reflection_servicer_test.py

@@ -17,15 +17,9 @@ import unittest
 
 import grpc
 
-# TODO(https://github.com/grpc/grpc/issues/19863): Remove.
-try:
-    from src.python.grpcio_reflection.grpc_reflection.v1alpha import reflection
-    from src.python.grpcio_reflection.grpc_reflection.v1alpha import reflection_pb2
-    from src.python.grpcio_reflection.grpc_reflection.v1alpha import reflection_pb2_grpc
-except ImportError:
-    from grpc_reflection.v1alpha import reflection
-    from grpc_reflection.v1alpha import reflection_pb2
-    from grpc_reflection.v1alpha import reflection_pb2_grpc
+from grpc_reflection.v1alpha import reflection
+from grpc_reflection.v1alpha import reflection_pb2
+from grpc_reflection.v1alpha import reflection_pb2_grpc
 
 from google.protobuf import descriptor_pool
 from google.protobuf import descriptor_pb2