Make Python package spec indirect

This is part of a change to ease internal usage of GRPC.

1 files changed, 60 insertions, 9 deletions

diff --git a/src/compiler/python_generator.cc b/src/compiler/python_generator.cc
index e4f85450f5..1ec3772f9f 100644
--- a/src/compiler/python_generator.cc
+++ b/src/compiler/python_generator.cc
@@ -36,13 +36,18 @@
 #include <cctype>
 #include <cstring>
 #include <map>
+#include <memory>
 #include <ostream>
 #include <sstream>
+#include <string>
+#include <tuple>
 #include <vector>
 
 #include "src/compiler/generator_helpers.h"
 #include "src/compiler/python_generator.h"
+#include <google/protobuf/io/coded_stream.h>
 #include <google/protobuf/io/printer.h>
+#include <google/protobuf/io/zero_copy_stream.h>
 #include <google/protobuf/io/zero_copy_stream_impl_lite.h>
 #include <google/protobuf/descriptor.pb.h>
 #include <google/protobuf/descriptor.h>
@@ -53,8 +58,11 @@ using google::protobuf::Descriptor;
 using google::protobuf::FileDescriptor;
 using google::protobuf::MethodDescriptor;
 using google::protobuf::ServiceDescriptor;
+using google::protobuf::compiler::GeneratorContext;
+using google::protobuf::io::CodedOutputStream;
 using google::protobuf::io::Printer;
 using google::protobuf::io::StringOutputStream;
+using google::protobuf::io::ZeroCopyOutputStream;
 using std::initializer_list;
 using std::make_pair;
 using std::map;
@@ -63,6 +71,41 @@ using std::replace;
 using std::vector;
 
 namespace grpc_python_generator {
+
+PythonGrpcGenerator::PythonGrpcGenerator(const GeneratorConfiguration& config)
+    : config_(config) {}
+
+PythonGrpcGenerator::~PythonGrpcGenerator() {}
+
+bool PythonGrpcGenerator::Generate(
+    const FileDescriptor* file, const std::string& parameter,
+    GeneratorContext* context, std::string* error) const {
+  // Get output file name.
+  std::string file_name;
+  static const int proto_suffix_length = strlen(".proto");
+  if (file->name().size() > static_cast<size_t>(proto_suffix_length) &&
+      file->name().find_last_of(".proto") == file->name().size() - 1) {
+    file_name = file->name().substr(
+        0, file->name().size() - proto_suffix_length) + "_pb2.py";
+  } else {
+    *error = "Invalid proto file name. Proto file must end with .proto";
+    return false;
+  }
+
+  std::unique_ptr<ZeroCopyOutputStream> output(
+      context->OpenForInsert(file_name, "module_scope"));
+  CodedOutputStream coded_out(output.get());
+  bool success = false;
+  std::string code = "";
+  tie(success, code) = grpc_python_generator::GetServices(file, config_);
+  if (success) {
+    coded_out.WriteRaw(code.data(), code.size());
+    return true;
+  } else {
+    return false;
+  }
+}
+
 namespace {
 //////////////////////////////////
 // BEGIN FORMATTING BOILERPLATE //
@@ -70,7 +113,8 @@ namespace {
 
 // Converts an initializer list of the form { key0, value0, key1, value1, ... }
 // into a map of key* to value*. Is merely a readability helper for later code.
-map<std::string, std::string> ListToDict(const initializer_list<std::string>& values) {
+map<std::string, std::string> ListToDict(
+    const initializer_list<std::string>& values) {
   assert(values.size() % 2 == 0);
   map<std::string, std::string> value_map;
   auto value_iter = values.begin();
@@ -237,8 +281,10 @@ bool PrintServerFactory(const std::string& package_qualified_service_name,
   {
     IndentScope raii_create_server_indent(out);
     map<std::string, std::string> method_description_constructors;
-    map<std::string, pair<std::string, std::string>> input_message_modules_and_classes;
-    map<std::string, pair<std::string, std::string>> output_message_modules_and_classes;
+    map<std::string, pair<std::string, std::string>>
+        input_message_modules_and_classes;
+    map<std::string, pair<std::string, std::string>>
+        output_message_modules_and_classes;
     for (int i = 0; i < service->method_count(); ++i) {
       const MethodDescriptor* method = service->method(i);
       const std::string method_description_constructor =
@@ -313,8 +359,10 @@ bool PrintStubFactory(const std::string& package_qualified_service_name,
   {
     IndentScope raii_create_server_indent(out);
     map<std::string, std::string> method_description_constructors;
-    map<std::string, pair<std::string, std::string>> input_message_modules_and_classes;
-    map<std::string, pair<std::string, std::string>> output_message_modules_and_classes;
+    map<std::string, pair<std::string, std::string>>
+        input_message_modules_and_classes;
+    map<std::string, pair<std::string, std::string>>
+        output_message_modules_and_classes;
     for (int i = 0; i < service->method_count(); ++i) {
       const MethodDescriptor* method = service->method(i);
       const std::string method_description_constructor =
@@ -378,22 +426,25 @@ bool PrintStubFactory(const std::string& package_qualified_service_name,
   return true;
 }
 
-bool PrintPreamble(const FileDescriptor* file, Printer* out) {
+bool PrintPreamble(const FileDescriptor* file,
+                   const GeneratorConfiguration& config, Printer* out) {
   out->Print("import abc\n");
-  out->Print("from grpc.early_adopter import implementations\n");
+  out->Print("from $Package$ import implementations\n",
+             "Package", config.implementations_package_root);
   out->Print("from grpc.framework.alpha import utilities\n");
   return true;
 }
 
 }  // namespace
 
-pair<bool, std::string> GetServices(const FileDescriptor* file) {
+pair<bool, std::string> GetServices(const FileDescriptor* file,
+                                    const GeneratorConfiguration& config) {
   std::string output;
   {
     // Scope the output stream so it closes and finalizes output to the string.
     StringOutputStream output_stream(&output);
     Printer out(&output_stream, '$');
-    if (!PrintPreamble(file, &out)) {
+    if (!PrintPreamble(file, config, &out)) {
       return make_pair(false, "");
     }
     auto package = file->package();