Merge pull request #18555 from yang-g/proto_h

Add options to support alternative extension for include files

src/compiler/cpp_generator.cc

@@ -84,7 +84,7 @@ void PrintIncludes(grpc_generator::Printer* printer,
 }
 
 grpc::string GetHeaderPrologue(grpc_generator::File* file,
-                               const Parameters& /*params*/) {
+                               const Parameters& params) {
   grpc::string output;
   {
     // Scope the output stream so it closes and finalizes output to the string.
@@ -94,7 +94,9 @@ grpc::string GetHeaderPrologue(grpc_generator::File* file,
     vars["filename"] = file->filename();
     vars["filename_identifier"] = FilenameIdentifier(file->filename());
     vars["filename_base"] = file->filename_without_ext();
-    vars["message_header_ext"] = kCppGeneratorMessageHeaderExt;
+    vars["message_header_ext"] = params.message_header_extension.empty()
+                                     ? kCppGeneratorMessageHeaderExt
+                                     : params.message_header_extension;
 
     printer->Print(vars, "// Generated by the gRPC C++ plugin.\n");
     printer->Print(vars,
@@ -115,6 +117,13 @@ grpc::string GetHeaderPrologue(grpc_generator::File* file,
   return output;
 }
 
+// Convert from "a/b/c.proto" to "#include \"a/b/c$message_header_ext$\"\n"
+grpc::string ImportInludeFromProtoName(const grpc::string& proto_name) {
+  return grpc::string("#include \"") +
+         proto_name.substr(0, proto_name.size() - 6) +
+         grpc::string("$message_header_ext$\"\n");
+}
+
 grpc::string GetHeaderIncludes(grpc_generator::File* file,
                                const Parameters& params) {
   grpc::string output;
@@ -152,6 +161,20 @@ grpc::string GetHeaderIncludes(grpc_generator::File* file,
     printer->Print(vars, "class ServerContext;\n");
     printer->Print(vars, "}  // namespace grpc\n\n");
 
+    vars["message_header_ext"] = params.message_header_extension.empty()
+                                     ? kCppGeneratorMessageHeaderExt
+                                     : params.message_header_extension;
+
+    if (params.include_import_headers) {
+      const std::vector<grpc::string> import_names = file->GetImportNames();
+      for (const auto& import_name : import_names) {
+        const grpc::string include_name =
+            ImportInludeFromProtoName(import_name);
+        printer->Print(vars, include_name.c_str());
+      }
+      printer->PrintRaw("\n");
+    }
+
     if (!file->package().empty()) {
       std::vector<grpc::string> parts = file->package_parts();
 
@@ -1556,7 +1579,7 @@ grpc::string GetHeaderEpilogue(grpc_generator::File* file,
 }
 
 grpc::string GetSourcePrologue(grpc_generator::File* file,
-                               const Parameters& /*params*/) {
+                               const Parameters& params) {
   grpc::string output;
   {
     // Scope the output stream so it closes and finalizes output to the string.
@@ -1565,7 +1588,9 @@ grpc::string GetSourcePrologue(grpc_generator::File* file,
 
     vars["filename"] = file->filename();
     vars["filename_base"] = file->filename_without_ext();
-    vars["message_header_ext"] = kCppGeneratorMessageHeaderExt;
+    vars["message_header_ext"] = params.message_header_extension.empty()
+                                     ? kCppGeneratorMessageHeaderExt
+                                     : params.message_header_extension;
     vars["service_header_ext"] = kCppGeneratorServiceHeaderExt;
 
     printer->Print(vars, "// Generated by the gRPC C++ plugin.\n");
@@ -1587,7 +1612,6 @@ grpc::string GetSourceIncludes(grpc_generator::File* file,
     // Scope the output stream so it closes and finalizes output to the string.
     auto printer = file->CreatePrinter(&output);
     std::map<grpc::string, grpc::string> vars;
-
     static const char* headers_strs[] = {
         "functional",
         "grpcpp/impl/codegen/async_stream.h",
@@ -2046,7 +2070,7 @@ grpc::string GetSourceEpilogue(grpc_generator::File* file,
 
 // TODO(mmukhi): Make sure we need parameters or not.
 grpc::string GetMockPrologue(grpc_generator::File* file,
-                             const Parameters& /*params*/) {
+                             const Parameters& params) {
   grpc::string output;
   {
     // Scope the output stream so it closes and finalizes output to the string.
@@ -2055,7 +2079,9 @@ grpc::string GetMockPrologue(grpc_generator::File* file,
 
     vars["filename"] = file->filename();
     vars["filename_base"] = file->filename_without_ext();
-    vars["message_header_ext"] = kCppGeneratorMessageHeaderExt;
+    vars["message_header_ext"] = params.message_header_extension.empty()
+                                     ? kCppGeneratorMessageHeaderExt
+                                     : params.message_header_extension;
     vars["service_header_ext"] = kCppGeneratorServiceHeaderExt;
 
     printer->Print(vars, "// Generated by the gRPC C++ plugin.\n");
@@ -2065,6 +2091,15 @@ grpc::string GetMockPrologue(grpc_generator::File* file,
 
     printer->Print(vars, "#include \"$filename_base$$message_header_ext$\"\n");
     printer->Print(vars, "#include \"$filename_base$$service_header_ext$\"\n");
+    if (params.include_import_headers) {
+      const std::vector<grpc::string> import_names = file->GetImportNames();
+      for (const auto& import_name : import_names) {
+        const grpc::string include_name =
+            ImportInludeFromProtoName(import_name);
+        printer->Print(vars, include_name.c_str());
+      }
+      printer->PrintRaw("\n");
+    }
     printer->Print(vars, file->additional_headers().c_str());
     printer->Print(vars, "\n");
   }

src/compiler/cpp_generator.h

@@ -56,6 +56,10 @@ struct Parameters {
   grpc::string gmock_search_path;
   // *EXPERIMENTAL* Additional include files in grpc.pb.h
   std::vector<grpc::string> additional_header_includes;
+  // By default, use "pb.h"
+  grpc::string message_header_extension;
+  // Whether to include headers corresponding to imports in source file.
+  bool include_import_headers;
 };
 
 // Return the prologue of the generated header file.

src/compiler/cpp_plugin.cc

@@ -48,6 +48,7 @@ class CppGrpcGenerator : public grpc::protobuf::compiler::CodeGenerator {
     grpc_cpp_generator::Parameters generator_parameters;
     generator_parameters.use_system_headers = true;
     generator_parameters.generate_mock_code = false;
+    generator_parameters.include_import_headers = false;
 
     ProtoBufFile pbfile(file);
 
@@ -83,6 +84,15 @@ class CppGrpcGenerator : public grpc::protobuf::compiler::CodeGenerator {
         } else if (param[0] == "additional_header_includes") {
           generator_parameters.additional_header_includes =
               grpc_generator::tokenize(param[1], ":");
+        } else if (param[0] == "message_header_extension") {
+          generator_parameters.message_header_extension = param[1];
+        } else if (param[0] == "include_import_headers") {
+          if (param[1] == "true") {
+            generator_parameters.include_import_headers = true;
+          } else if (param[1] != "false") {
+            *error = grpc::string("Invalid parameter: ") + *parameter_string;
+            return false;
+          }
         } else {
           *error = grpc::string("Unknown parameter: ") + *parameter_string;
           return false;

src/compiler/protobuf_plugin.h

@@ -189,6 +189,15 @@ class ProtoBufFile : public grpc_generator::File {
     return grpc_python_generator::get_all_comments(file_);
   }
 
+  vector<grpc::string> GetImportNames() const {
+    vector<grpc::string> proto_names;
+    for (int i = 0; i < file_->dependency_count(); ++i) {
+      const auto& dep = *file_->dependency(i);
+      proto_names.push_back(dep.name());
+    }
+    return proto_names;
+  }
+
  private:
   const grpc::protobuf::FileDescriptor* file_;
 };

src/compiler/schema_interface.h

@@ -101,6 +101,7 @@ struct File : public CommentHolder {
   virtual grpc::string package() const = 0;
   virtual std::vector<grpc::string> package_parts() const = 0;
   virtual grpc::string additional_headers() const = 0;
+  virtual std::vector<grpc::string> GetImportNames() const { return {}; }
 
   virtual int service_count() const = 0;
   virtual std::unique_ptr<const Service> service(int i) const = 0;