Merge pull request #18808 from grpc/protoc_strip_module_prefix

Add ability to strip module prefixes in Python gRPC protoc plugin

src/compiler/protobuf_plugin.h

@@ -55,22 +55,23 @@ class ProtoBufMethod : public grpc_generator::Method {
     return method_->output_type()->file()->name();
   }
 
-  bool get_module_and_message_path_input(grpc::string* str,
+  // TODO(https://github.com/grpc/grpc/issues/18800): Clean this up.
-                                         grpc::string generator_file_name,
+  bool get_module_and_message_path_input(
-                                         bool generate_in_pb2_grpc,
+      grpc::string* str, grpc::string generator_file_name,
-                                         grpc::string import_prefix) const {
+      bool generate_in_pb2_grpc, grpc::string import_prefix,
+      const std::vector<grpc::string>& prefixes_to_filter) const final {
     return grpc_python_generator::GetModuleAndMessagePath(
         method_->input_type(), str, generator_file_name, generate_in_pb2_grpc,
-        import_prefix);
+        import_prefix, prefixes_to_filter);
   }
 
-  bool get_module_and_message_path_output(grpc::string* str,
+  bool get_module_and_message_path_output(
-                                          grpc::string generator_file_name,
+      grpc::string* str, grpc::string generator_file_name,
-                                          bool generate_in_pb2_grpc,
+      bool generate_in_pb2_grpc, grpc::string import_prefix,
-                                          grpc::string import_prefix) const {
+      const std::vector<grpc::string>& prefixes_to_filter) const final {
     return grpc_python_generator::GetModuleAndMessagePath(
         method_->output_type(), str, generator_file_name, generate_in_pb2_grpc,
-        import_prefix);
+        import_prefix, prefixes_to_filter);
   }
 
   bool NoStreaming() const {

src/compiler/python_generator.cc

@@ -214,13 +214,15 @@ bool PrivateGenerator::PrintBetaServerFactory(
       grpc::string input_message_module_and_class;
       if (!method->get_module_and_message_path_input(
               &input_message_module_and_class, generator_file_name,
-              generate_in_pb2_grpc, config.import_prefix)) {
+              generate_in_pb2_grpc, config.import_prefix,
+              config.prefixes_to_filter)) {
         return false;
       }
       grpc::string output_message_module_and_class;
       if (!method->get_module_and_message_path_output(
               &output_message_module_and_class, generator_file_name,
-              generate_in_pb2_grpc, config.import_prefix)) {
+              generate_in_pb2_grpc, config.import_prefix,
+              config.prefixes_to_filter)) {
         return false;
       }
       method_implementation_constructors.insert(
@@ -320,13 +322,15 @@ bool PrivateGenerator::PrintBetaStubFactory(
       grpc::string input_message_module_and_class;
       if (!method->get_module_and_message_path_input(
               &input_message_module_and_class, generator_file_name,
-              generate_in_pb2_grpc, config.import_prefix)) {
+              generate_in_pb2_grpc, config.import_prefix,
+              config.prefixes_to_filter)) {
         return false;
       }
       grpc::string output_message_module_and_class;
       if (!method->get_module_and_message_path_output(
               &output_message_module_and_class, generator_file_name,
-              generate_in_pb2_grpc, config.import_prefix)) {
+              generate_in_pb2_grpc, config.import_prefix,
+              config.prefixes_to_filter)) {
         return false;
       }
       method_cardinalities.insert(
@@ -425,13 +429,15 @@ bool PrivateGenerator::PrintStub(
         grpc::string request_module_and_class;
         if (!method->get_module_and_message_path_input(
                 &request_module_and_class, generator_file_name,
-                generate_in_pb2_grpc, config.import_prefix)) {
+                generate_in_pb2_grpc, config.import_prefix,
+                config.prefixes_to_filter)) {
           return false;
         }
         grpc::string response_module_and_class;
         if (!method->get_module_and_message_path_output(
                 &response_module_and_class, generator_file_name,
-                generate_in_pb2_grpc, config.import_prefix)) {
+                generate_in_pb2_grpc, config.import_prefix,
+                config.prefixes_to_filter)) {
           return false;
         }
         StringMap method_dict;
@@ -516,13 +522,15 @@ bool PrivateGenerator::PrintAddServicerToServer(
         grpc::string request_module_and_class;
         if (!method->get_module_and_message_path_input(
                 &request_module_and_class, generator_file_name,
-                generate_in_pb2_grpc, config.import_prefix)) {
+                generate_in_pb2_grpc, config.import_prefix,
+                config.prefixes_to_filter)) {
           return false;
         }
         grpc::string response_module_and_class;
         if (!method->get_module_and_message_path_output(
                 &response_module_and_class, generator_file_name,
-                generate_in_pb2_grpc, config.import_prefix)) {
+                generate_in_pb2_grpc, config.import_prefix,
+                config.prefixes_to_filter)) {
           return false;
         }
         StringMap method_dict;
@@ -589,17 +597,21 @@ bool PrivateGenerator::PrintPreamble(grpc_generator::Printer* out) {
 
         grpc::string input_type_file_name = method->get_input_type_name();
         grpc::string input_module_name =
-            ModuleName(input_type_file_name, config.import_prefix);
+            ModuleName(input_type_file_name, config.import_prefix,
+                       config.prefixes_to_filter);
         grpc::string input_module_alias =
-            ModuleAlias(input_type_file_name, config.import_prefix);
+            ModuleAlias(input_type_file_name, config.import_prefix,
+                        config.prefixes_to_filter);
         imports_set.insert(
             std::make_tuple(input_module_name, input_module_alias));
 
         grpc::string output_type_file_name = method->get_output_type_name();
         grpc::string output_module_name =
-            ModuleName(output_type_file_name, config.import_prefix);
+            ModuleName(output_type_file_name, config.import_prefix,
+                       config.prefixes_to_filter);
         grpc::string output_module_alias =
-            ModuleAlias(output_type_file_name, config.import_prefix);
+            ModuleAlias(output_type_file_name, config.import_prefix,
+                        config.prefixes_to_filter);
         imports_set.insert(
             std::make_tuple(output_module_name, output_module_alias));
       }

src/compiler/python_generator.h

@@ -20,6 +20,7 @@
 #define GRPC_INTERNAL_COMPILER_PYTHON_GENERATOR_H
 
 #include <utility>
+#include <vector>
 
 #include "src/compiler/config.h"
 #include "src/compiler/schema_interface.h"
@@ -35,6 +36,7 @@ struct GeneratorConfiguration {
   grpc::string beta_package_root;
   // TODO(https://github.com/google/protobuf/issues/888): Drop this.
   grpc::string import_prefix;
+  std::vector<grpc::string> prefixes_to_filter;
 };
 
 class PythonGrpcGenerator : public grpc::protobuf::compiler::CodeGenerator {

src/compiler/python_generator_helpers.h

@@ -49,23 +49,39 @@ namespace {
 typedef vector<const Descriptor*> DescriptorVector;
 typedef vector<grpc::string> StringVector;
 
+static grpc::string StripModulePrefixes(
+    const grpc::string& raw_module_name,
+    const std::vector<grpc::string>& prefixes_to_filter) {
+  for (const auto& prefix : prefixes_to_filter) {
+    if (raw_module_name.rfind(prefix, 0) == 0) {
+      return raw_module_name.substr(prefix.size(),
+                                    raw_module_name.size() - prefix.size());
+    }
+  }
+  return raw_module_name;
+}
+
 // TODO(https://github.com/google/protobuf/issues/888):
 // Export `ModuleName` from protobuf's
 // `src/google/protobuf/compiler/python/python_generator.cc` file.
 grpc::string ModuleName(const grpc::string& filename,
-                        const grpc::string& import_prefix) {
+                        const grpc::string& import_prefix,
+                        const std::vector<grpc::string>& prefixes_to_filter) {
   grpc::string basename = StripProto(filename);
   basename = StringReplace(basename, "-", "_");
   basename = StringReplace(basename, "/", ".");
-  return import_prefix + basename + "_pb2";
+  return StripModulePrefixes(import_prefix + basename + "_pb2",
+                             prefixes_to_filter);
 }
 
 // TODO(https://github.com/google/protobuf/issues/888):
 // Export `ModuleAlias` from protobuf's
 // `src/google/protobuf/compiler/python/python_generator.cc` file.
 grpc::string ModuleAlias(const grpc::string& filename,
-                         const grpc::string& import_prefix) {
+                         const grpc::string& import_prefix,
-  grpc::string module_name = ModuleName(filename, import_prefix);
+                         const std::vector<grpc::string>& prefixes_to_filter) {
+  grpc::string module_name =
+      ModuleName(filename, import_prefix, prefixes_to_filter);
   // We can't have dots in the module name, so we replace each with _dot_.
   // But that could lead to a collision between a.b and a_dot_b, so we also
   // duplicate each underscore.
@@ -74,10 +90,10 @@ grpc::string ModuleAlias(const grpc::string& filename,
   return module_name;
 }
 
-bool GetModuleAndMessagePath(const Descriptor* type, grpc::string* out,
+bool GetModuleAndMessagePath(
-                             grpc::string generator_file_name,
+    const Descriptor* type, grpc::string* out, grpc::string generator_file_name,
-                             bool generate_in_pb2_grpc,
+    bool generate_in_pb2_grpc, grpc::string& import_prefix,
-                             grpc::string& import_prefix) {
+    const std::vector<grpc::string>& prefixes_to_filter) {
   const Descriptor* path_elem_type = type;
   DescriptorVector message_path;
   do {
@@ -93,7 +109,7 @@ bool GetModuleAndMessagePath(const Descriptor* type, grpc::string* out,
 
   grpc::string module;
   if (generator_file_name != file_name || generate_in_pb2_grpc) {
-    module = ModuleAlias(file_name, import_prefix) + ".";
+    module = ModuleAlias(file_name, import_prefix, prefixes_to_filter) + ".";
   } else {
     module = "";
   }

src/compiler/schema_interface.h

@@ -57,10 +57,12 @@ struct Method : public CommentHolder {
 
   virtual bool get_module_and_message_path_input(
       grpc::string* str, grpc::string generator_file_name,
-      bool generate_in_pb2_grpc, grpc::string import_prefix) const = 0;
+      bool generate_in_pb2_grpc, grpc::string import_prefix,
+      const std::vector<grpc::string>& prefixes_to_filter) const = 0;
   virtual bool get_module_and_message_path_output(
       grpc::string* str, grpc::string generator_file_name,
-      bool generate_in_pb2_grpc, grpc::string import_prefix) const = 0;
+      bool generate_in_pb2_grpc, grpc::string import_prefix,
+      const std::vector<grpc::string>& prefixes_to_filter) const = 0;
 
   virtual grpc::string get_input_type_name() const = 0;
   virtual grpc::string get_output_type_name() const = 0;