Merge pull request #4679 from nicolasnoble/codegen-paths

Letting the user override the code generation a bit.

3 files changed, 89 insertions, 39 deletions

diff --git a/src/compiler/cpp_generator.cc b/src/compiler/cpp_generator.cc
index 97295bed45..206a6e1fe5 100644
--- a/src/compiler/cpp_generator.cc
+++ b/src/compiler/cpp_generator.cc
@@ -83,6 +83,28 @@ grpc::string FilenameIdentifier(const grpc::string &filename) {
 }
 }  // namespace
 
+template<class T, size_t N>
+T *array_end(T (&array)[N]) { return array + N; }
+
+void PrintIncludes(grpc::protobuf::io::Printer *printer, const std::vector<grpc::string>& headers, const Parameters &params) {
+  std::map<grpc::string, grpc::string> vars;
+
+  vars["l"] = params.use_system_headers ? '<' : '"';
+  vars["r"] = params.use_system_headers ? '>' : '"';
+
+  if (!params.grpc_search_path.empty()) {
+    vars["l"] += params.grpc_search_path;
+    if (params.grpc_search_path.back() != '/') {
+      vars["l"] += '/';
+    }
+  }
+
+  for (auto i = headers.begin(); i != headers.end(); i++) {
+    vars["h"] = *i;
+    printer->Print(vars, "#include $l$$h$$r$\n");
+  }
+}
+
 grpc::string GetHeaderPrologue(const grpc::protobuf::FileDescriptor *file,
                                const Parameters &params) {
   grpc::string output;
@@ -111,36 +133,46 @@ grpc::string GetHeaderPrologue(const grpc::protobuf::FileDescriptor *file,
 
 grpc::string GetHeaderIncludes(const grpc::protobuf::FileDescriptor *file,
                                const Parameters &params) {
-  grpc::string temp =
-      "#include <grpc++/impl/codegen/async_stream.h>\n"
-      "#include <grpc++/impl/codegen/async_unary_call.h>\n"
-      "#include <grpc++/impl/codegen/proto_utils.h>\n"
-      "#include <grpc++/impl/codegen/rpc_method.h>\n"
-      "#include <grpc++/impl/codegen/service_type.h>\n"
-      "#include <grpc++/impl/codegen/status.h>\n"
-      "#include <grpc++/impl/codegen/stub_options.h>\n"
-      "#include <grpc++/impl/codegen/sync_stream.h>\n"
-      "\n"
-      "namespace grpc {\n"
-      "class CompletionQueue;\n"
-      "class RpcService;\n"
-      "class ServerCompletionQueue;\n"
-      "class ServerContext;\n"
-      "}  // namespace grpc\n\n";
+  grpc::string output;
+  {
+    // Scope the output stream so it closes and finalizes output to the string.
+    grpc::protobuf::io::StringOutputStream output_stream(&output);
+    grpc::protobuf::io::Printer printer(&output_stream, '$');
+    std::map<grpc::string, grpc::string> vars;
 
-  if (!file->package().empty()) {
-    std::vector<grpc::string> parts =
-        grpc_generator::tokenize(file->package(), ".");
+    static const char *headers_strs[] = {
+      "grpc++/impl/codegen/async_stream.h",
+      "grpc++/impl/codegen/async_unary_call.h",
+      "grpc++/impl/codegen/proto_utils.h",
+      "grpc++/impl/codegen/rpc_method.h",
+      "grpc++/impl/codegen/service_type.h",
+      "grpc++/impl/codegen/status.h",
+      "grpc++/impl/codegen/stub_options.h",
+      "grpc++/impl/codegen/sync_stream.h"
+    };
+    std::vector<grpc::string> headers(headers_strs, array_end(headers_strs));
+    PrintIncludes(&printer, headers, params);
+    printer.Print(vars, "\n");
+    printer.Print(vars, "namespace grpc {\n");
+    printer.Print(vars, "class CompletionQueue;\n");
+    printer.Print(vars, "class Channel;\n");
+    printer.Print(vars, "class RpcService;\n");
+    printer.Print(vars, "class ServerCompletionQueue;\n");
+    printer.Print(vars, "class ServerContext;\n");
+    printer.Print(vars, "}  // namespace grpc\n\n");
 
-    for (auto part = parts.begin(); part != parts.end(); part++) {
-      temp.append("namespace ");
-      temp.append(*part);
-      temp.append(" {\n");
+    if (!file->package().empty()) {
+      std::vector<grpc::string> parts =
+          grpc_generator::tokenize(file->package(), ".");
+
+      for (auto part = parts.begin(); part != parts.end(); part++) {
+        vars["part"] = *part;
+        printer.Print(vars, "namespace $part$ {\n");
+      }
+      printer.Print(vars, "\n");
     }
-    temp.append("\n");
   }
-
-  return temp;
+  return output;
 }
 
 void PrintHeaderClientMethodInterfaces(
@@ -852,7 +884,7 @@ grpc::string GetSourcePrologue(const grpc::protobuf::FileDescriptor *file,
 }
 
 grpc::string GetSourceIncludes(const grpc::protobuf::FileDescriptor *file,
-                               const Parameters &param) {
+                               const Parameters &params) {
   grpc::string output;
   {
     // Scope the output stream so it closes and finalizes output to the string.
@@ -860,16 +892,18 @@ grpc::string GetSourceIncludes(const grpc::protobuf::FileDescriptor *file,
     grpc::protobuf::io::Printer printer(&output_stream, '$');
     std::map<grpc::string, grpc::string> vars;
 
-    printer.Print(vars, "#include <grpc++/impl/codegen/async_stream.h>\n");
-    printer.Print(vars, "#include <grpc++/impl/codegen/async_unary_call.h>\n");
-    printer.Print(vars, "#include <grpc++/impl/codegen/channel_interface.h>\n");
-    printer.Print(vars, "#include <grpc++/impl/codegen/client_unary_call.h>\n");
-    printer.Print(vars,
-                  "#include <grpc++/impl/codegen/method_handler_impl.h>\n");
-    printer.Print(vars,
-                  "#include <grpc++/impl/codegen/rpc_service_method.h>\n");
-    printer.Print(vars, "#include <grpc++/impl/codegen/service_type.h>\n");
-    printer.Print(vars, "#include <grpc++/impl/codegen/sync_stream.h>\n");
+    static const char *headers_strs[] = {
+      "grpc++/impl/codegen/async_stream.h",
+      "grpc++/impl/codegen/async_unary_call.h",
+      "grpc++/impl/codegen/channel_interface.h",
+      "grpc++/impl/codegen/client_unary_call.h",
+      "grpc++/impl/codegen/method_handler_impl.h",
+      "grpc++/impl/codegen/rpc_service_method.h",
+      "grpc++/impl/codegen/service_type.h",
+      "grpc++/impl/codegen/sync_stream.h"
+    };
+    std::vector<grpc::string> headers(headers_strs, array_end(headers_strs));
+    PrintIncludes(&printer, headers, params);
 
     if (!file->package().empty()) {
       std::vector<grpc::string> parts =
diff --git a/src/compiler/cpp_generator.h b/src/compiler/cpp_generator.h
index 70c2e985f6..4f9de9d11a 100644
--- a/src/compiler/cpp_generator.h
+++ b/src/compiler/cpp_generator.h
@@ -1,6 +1,6 @@
 /*
  *
- * Copyright 2015, Google Inc.
+ * Copyright 2015-2016, Google Inc.
  * All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
@@ -42,6 +42,10 @@ namespace grpc_cpp_generator {
 struct Parameters {
   // Puts the service into a namespace
   grpc::string services_namespace;
+  // Use system includes (<>) or local includes ("")
+  bool use_system_headers;
+  // Prefix to any grpc include
+  grpc::string grpc_search_path;
 };
 
 // Return the prologue of the generated header file.
diff --git a/src/compiler/cpp_plugin.cc b/src/compiler/cpp_plugin.cc
index 88c704948e..d8ada4835c 100644
--- a/src/compiler/cpp_plugin.cc
+++ b/src/compiler/cpp_plugin.cc
@@ -1,6 +1,6 @@
 /*
  *
- * Copyright 2015, Google Inc.
+ * Copyright 2015-2016, Google Inc.
  * All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
@@ -59,6 +59,7 @@ class CppGrpcGenerator : public grpc::protobuf::compiler::CodeGenerator {
     }
 
     grpc_cpp_generator::Parameters generator_parameters;
+    generator_parameters.use_system_headers = true;
 
     if (!parameter.empty()) {
       std::vector<grpc::string> parameters_list =
@@ -70,6 +71,17 @@ class CppGrpcGenerator : public grpc::protobuf::compiler::CodeGenerator {
           grpc_generator::tokenize(*parameter_string, "=");
         if (param[0] == "services_namespace") {
           generator_parameters.services_namespace = param[1];
+        } else if (param[0] == "use_system_headers") {
+          if (param[1] == "true") {
+            generator_parameters.use_system_headers = true;
+          } else if (param[1] == "false") {
+            generator_parameters.use_system_headers = false;
+          } else {
+            *error = grpc::string("Invalid parameter: ") + *parameter_string;
+            return false;
+          }
+        } else if (param[0] == "grpc_search_path") {
+          generator_parameters.grpc_search_path = param[1];
         } else {
           *error = grpc::string("Unknown parameter: ") + *parameter_string;
           return false;