1 files changed, 49 insertions, 23 deletions

diff --git a/src/compiler/python_generator.cc b/src/compiler/python_generator.cc
index febaf135b6..4841da8da8 100644
--- a/src/compiler/python_generator.cc
+++ b/src/compiler/python_generator.cc
@@ -40,6 +40,7 @@
 #include <map>
 #include <memory>
 #include <ostream>
+#include <set>
 #include <sstream>
 #include <tuple>
 #include <vector>
@@ -64,7 +65,9 @@ using std::make_pair;
 using std::map;
 using std::pair;
 using std::replace;
+using std::tuple;
 using std::vector;
+using std::set;
 
 namespace grpc_python_generator {
 
@@ -73,6 +76,8 @@ namespace {
 typedef vector<const Descriptor*> DescriptorVector;
 typedef map<grpc::string, grpc::string> StringMap;
 typedef vector<grpc::string> StringVector;
+typedef tuple<grpc::string, grpc::string> StringPair;
+typedef set<StringPair> StringPairSet;
 
 // Provides RAII indentation handling. Use as:
 // {
@@ -651,6 +656,7 @@ bool PrivateGenerator::PrintPreamble() {
       "face_utilities\n");
   if (generate_in_pb2_grpc) {
     out->Print("\n");
+    StringPairSet imports_set;
     for (int i = 0; i < file->service_count(); ++i) {
       const ServiceDescriptor* service = file->service(i);
       for (int j = 0; j < service->method_count(); ++j) {
@@ -662,11 +668,15 @@ bool PrivateGenerator::PrintPreamble() {
           grpc::string type_file_name = type->file()->name();
           grpc::string module_name = ModuleName(type_file_name);
           grpc::string module_alias = ModuleAlias(type_file_name);
-          out->Print("import $ModuleName$ as $ModuleAlias$\n", "ModuleName",
-                     module_name, "ModuleAlias", module_alias);
+          imports_set.insert(std::make_tuple(module_name, module_alias));
         }
       }
     }
+    for (StringPairSet::iterator it = imports_set.begin();
+         it != imports_set.end(); ++it) {
+      out->Print("import $ModuleName$ as $ModuleAlias$\n", "ModuleName",
+                 std::get<0>(*it), "ModuleAlias", std::get<1>(*it));
+    }
   }
   return true;
 }
@@ -714,6 +724,9 @@ pair<bool, grpc::string> PrivateGenerator::GetGrpcServices() {
     out = &out_printer;
 
     if (generate_in_pb2_grpc) {
+      out->Print(
+          "# Generated by the gRPC Python protocol compiler plugin. "
+          "DO NOT EDIT!\n");
       if (!PrintPreamble()) {
         return make_pair(false, "");
       }
@@ -760,6 +773,32 @@ PythonGrpcGenerator::PythonGrpcGenerator(const GeneratorConfiguration& config)
 
 PythonGrpcGenerator::~PythonGrpcGenerator() {}
 
+static bool GenerateGrpc(GeneratorContext* context, PrivateGenerator& generator,
+                         grpc::string file_name, bool generate_in_pb2_grpc) {
+  bool success;
+  std::unique_ptr<ZeroCopyOutputStream> output;
+  std::unique_ptr<CodedOutputStream> coded_output;
+  grpc::string grpc_code;
+
+  if (generate_in_pb2_grpc) {
+    output.reset(context->Open(file_name));
+    generator.generate_in_pb2_grpc = true;
+  } else {
+    output.reset(context->OpenForInsert(file_name, "module_scope"));
+    generator.generate_in_pb2_grpc = false;
+  }
+
+  coded_output.reset(new CodedOutputStream(output.get()));
+  tie(success, grpc_code) = generator.GetGrpcServices();
+
+  if (success) {
+    coded_output->WriteRaw(grpc_code.data(), grpc_code.size());
+    return true;
+  } else {
+    return false;
+  }
+}
+
 bool PythonGrpcGenerator::Generate(const FileDescriptor* file,
                                    const grpc::string& parameter,
                                    GeneratorContext* context,
@@ -780,28 +819,15 @@ bool PythonGrpcGenerator::Generate(const FileDescriptor* file,
   }
 
   PrivateGenerator generator(config_, file);
-
-  std::unique_ptr<ZeroCopyOutputStream> pb2_output(
-      context->OpenForAppend(pb2_file_name));
-  std::unique_ptr<ZeroCopyOutputStream> grpc_output(
-      context->Open(pb2_grpc_file_name));
-  CodedOutputStream pb2_coded_out(pb2_output.get());
-  CodedOutputStream grpc_coded_out(grpc_output.get());
-  bool success = false;
-  grpc::string pb2_code;
-  grpc::string grpc_code;
-  generator.generate_in_pb2_grpc = false;
-  tie(success, pb2_code) = generator.GetGrpcServices();
-  if (success) {
-    generator.generate_in_pb2_grpc = true;
-    tie(success, grpc_code) = generator.GetGrpcServices();
-    if (success) {
-      pb2_coded_out.WriteRaw(pb2_code.data(), pb2_code.size());
-      grpc_coded_out.WriteRaw(grpc_code.data(), grpc_code.size());
-      return true;
-    }
+  if (parameter == "grpc_2_0") {
+    return GenerateGrpc(context, generator, pb2_grpc_file_name, true);
+  } else if (parameter == "") {
+    return GenerateGrpc(context, generator, pb2_grpc_file_name, true) &&
+           GenerateGrpc(context, generator, pb2_file_name, false);
+  } else {
+    *error = "Invalid parameter '" + parameter + "'.";
+    return false;
   }
-  return false;
 }
 
 }  // namespace grpc_python_generator