From 561dce452524bfa13032a84523972acb5fcb180d Mon Sep 17 00:00:00 2001
From: Masood Malekghassemi <atash@google.com>
Date: Wed, 30 Nov 2016 20:20:32 -0800
Subject: [PATCH] Switch to command-arg for Python split-codegen

Almost fixes #8753 (which still needs a test).
---
 src/compiler/python_generator.cc              | 55 ++++++++++++-------
 .../protoc_plugin/_split_definitions_test.py  |  6 +-
 2 files changed, 37 insertions(+), 24 deletions(-)

diff --git a/src/compiler/python_generator.cc b/src/compiler/python_generator.cc
index febaf135b6f..b0a60092ab1 100644
--- a/src/compiler/python_generator.cc
+++ b/src/compiler/python_generator.cc
@@ -760,6 +760,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 +806,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
diff --git a/src/python/grpcio_tests/tests/protoc_plugin/_split_definitions_test.py b/src/python/grpcio_tests/tests/protoc_plugin/_split_definitions_test.py
index 089366a8c72..64fd97256eb 100644
--- a/src/python/grpcio_tests/tests/protoc_plugin/_split_definitions_test.py
+++ b/src/python/grpcio_tests/tests/protoc_plugin/_split_definitions_test.py
@@ -167,7 +167,7 @@ class SameSeparateTest(unittest.TestCase, SeparateTestMixin):
         '',
         '--proto_path={}'.format(self.proto_directory),
         '--python_out={}'.format(self.python_out_directory),
-        '--grpc_python_out={}'.format(self.grpc_python_out_directory),
+        '--grpc_python_out=grpc_2_0:{}'.format(self.grpc_python_out_directory),
         same_proto_file,
     ])
     if protoc_result != 0:
@@ -241,7 +241,7 @@ class SplitCommonTest(unittest.TestCase, CommonTestMixin):
         '',
         '--proto_path={}'.format(self.proto_directory),
         '--python_out={}'.format(self.python_out_directory),
-        '--grpc_python_out={}'.format(self.python_out_directory),
+        '--grpc_python_out={}'.format(self.grpc_python_out_directory),
         services_proto_file,
         messages_proto_file,
     ])
@@ -285,7 +285,7 @@ class SplitSeparateTest(unittest.TestCase, SeparateTestMixin):
         '',
         '--proto_path={}'.format(self.proto_directory),
         '--python_out={}'.format(self.python_out_directory),
-        '--grpc_python_out={}'.format(self.grpc_python_out_directory),
+        '--grpc_python_out=grpc_2_0:{}'.format(self.grpc_python_out_directory),
         services_proto_file,
         messages_proto_file,
     ])