Update Python generator files

1 files changed, 185 insertions, 269 deletions

diff --git a/src/compiler/python_generator.cc b/src/compiler/python_generator.cc
index 49d90fd36e..ff4a83e2a4 100644
--- a/src/compiler/python_generator.cc
+++ b/src/compiler/python_generator.cc
@@ -46,8 +46,12 @@
 #include <vector>
 
 #include "src/compiler/config.h"
+#include "src/compiler/schema_interface.h"
 #include "src/compiler/generator_helpers.h"
+#include "src/compiler/protobuf_plugin.h"
 #include "src/compiler/python_generator.h"
+#include "src/compiler/python_private_generator.h"
+#include "src/compiler/python_generator_helpers.h"
 
 using grpc_generator::StringReplace;
 using grpc_generator::StripProto;
@@ -79,6 +83,8 @@ typedef vector<grpc::string> StringVector;
 typedef tuple<grpc::string, grpc::string> StringPair;
 typedef set<StringPair> StringPairSet;
 
+grpc::string generator_file_name;
+
 // Provides RAII indentation handling. Use as:
 // {
 //   IndentScope raii_my_indent_var_name_here(my_py_printer);
@@ -88,152 +94,42 @@ typedef set<StringPair> StringPairSet;
 // }
 class IndentScope {
  public:
-  explicit IndentScope(Printer* printer) : printer_(printer) {
+  explicit IndentScope(grpc_generator::Printer* printer) : printer_(printer) {
     printer_->Indent();
   }
 
   ~IndentScope() { printer_->Outdent(); }
 
  private:
-  Printer* printer_;
-};
-
-// TODO(https://github.com/google/protobuf/issues/888):
-// Export `ModuleName` from protobuf's
-// `src/google/protobuf/compiler/python/python_generator.cc` file.
-grpc::string ModuleName(const grpc::string& filename,
-                        const grpc::string& import_prefix) {
-  grpc::string basename = StripProto(filename);
-  basename = StringReplace(basename, "-", "_");
-  basename = StringReplace(basename, "/", ".");
-  return import_prefix + basename + "_pb2";
-}
-
-// TODO(https://github.com/google/protobuf/issues/888):
-// Export `ModuleAlias` from protobuf's
-// `src/google/protobuf/compiler/python/python_generator.cc` file.
-grpc::string ModuleAlias(const grpc::string& filename,
-                         const grpc::string& import_prefix) {
-  grpc::string module_name = ModuleName(filename, import_prefix);
-  // We can't have dots in the module name, so we replace each with _dot_.
-  // But that could lead to a collision between a.b and a_dot_b, so we also
-  // duplicate each underscore.
-  module_name = StringReplace(module_name, "_", "__");
-  module_name = StringReplace(module_name, ".", "_dot_");
-  return module_name;
-}
-
-// Tucks all generator state in an anonymous namespace away from
-// PythonGrpcGenerator and the header file, mostly to encourage future changes
-// to not require updates to the grpcio-tools C++ code part. Assumes that it is
-// only ever used from a single thread.
-struct PrivateGenerator {
-  const GeneratorConfiguration& config;
-  const FileDescriptor* file;
-
-  bool generate_in_pb2_grpc;
-
-  Printer* out;
-
-  PrivateGenerator(const GeneratorConfiguration& config,
-                   const FileDescriptor* file);
-
-  std::pair<bool, grpc::string> GetGrpcServices();
-
- private:
-  bool PrintPreamble();
-  bool PrintBetaPreamble();
-  bool PrintGAServices();
-  bool PrintBetaServices();
-
-  bool PrintAddServicerToServer(
-      const grpc::string& package_qualified_service_name,
-      const ServiceDescriptor* service);
-  bool PrintServicer(const ServiceDescriptor* service);
-  bool PrintStub(const grpc::string& package_qualified_service_name,
-                 const ServiceDescriptor* service);
-
-  bool PrintBetaServicer(const ServiceDescriptor* service);
-  bool PrintBetaServerFactory(
-      const grpc::string& package_qualified_service_name,
-      const ServiceDescriptor* service);
-  bool PrintBetaStub(const ServiceDescriptor* service);
-  bool PrintBetaStubFactory(const grpc::string& package_qualified_service_name,
-                            const ServiceDescriptor* service);
-
-  // Get all comments (leading, leading_detached, trailing) and print them as a
-  // docstring. Any leading space of a line will be removed, but the line
-  // wrapping will not be changed.
-  template <typename DescriptorType>
-  void PrintAllComments(const DescriptorType* descriptor);
-
-  bool GetModuleAndMessagePath(const Descriptor* type, grpc::string* out);
+  grpc_generator::Printer* printer_;
 };
 
 PrivateGenerator::PrivateGenerator(const GeneratorConfiguration& config,
-                                   const FileDescriptor* file)
+                                   const grpc_generator::File* file)
     : config(config), file(file) {}
 
-bool PrivateGenerator::GetModuleAndMessagePath(const Descriptor* type,
-                                               grpc::string* out) {
-  const Descriptor* path_elem_type = type;
-  DescriptorVector message_path;
-  do {
-    message_path.push_back(path_elem_type);
-    path_elem_type = path_elem_type->containing_type();
-  } while (path_elem_type);  // implicit nullptr comparison; don't be explicit
-  grpc::string file_name = type->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)) {
-    return false;
-  }
-  grpc::string generator_file_name = file->name();
-  grpc::string module;
-  if (generator_file_name != file_name || generate_in_pb2_grpc) {
-    module = ModuleAlias(file_name, config.import_prefix) + ".";
-  } else {
-    module = "";
-  }
-  grpc::string message_type;
-  for (DescriptorVector::reverse_iterator path_iter = message_path.rbegin();
-       path_iter != message_path.rend(); ++path_iter) {
-    message_type += (*path_iter)->name() + ".";
-  }
-  // no pop_back prior to C++11
-  message_type.resize(message_type.size() - 1);
-  *out = module + message_type;
-  return true;
-}
+void PrivateGenerator::PrintAllComments(StringVector comments) {
 
-template <typename DescriptorType>
-void PrivateGenerator::PrintAllComments(const DescriptorType* descriptor) {
-  StringVector comments;
-  grpc_generator::GetComment(
-      descriptor, grpc_generator::COMMENTTYPE_LEADING_DETACHED, &comments);
-  grpc_generator::GetComment(descriptor, grpc_generator::COMMENTTYPE_LEADING,
-                             &comments);
-  grpc_generator::GetComment(descriptor, grpc_generator::COMMENTTYPE_TRAILING,
-                             &comments);
-  if (comments.empty()) {
-    return;
-  }
-  out->Print("\"\"\"");
-  for (StringVector::iterator it = comments.begin(); it != comments.end();
-       ++it) {
-    size_t start_pos = it->find_first_not_of(' ');
-    if (start_pos != grpc::string::npos) {
-      out->Print(it->c_str() + start_pos);
+    if (comments.empty()) {
+      return;
     }
-    out->Print("\n");
-  }
-  out->Print("\"\"\"\n");
+    out->Print("\"\"\"");
+    for (StringVector::iterator it = comments.begin(); it != comments.end();
+         ++it) {
+      size_t start_pos = it->find_first_not_of(' ');
+      if (start_pos != grpc::string::npos) {
+        out->Print(it->c_str() + start_pos);
+      }
+      out->Print("\n");
+    }
+    out->Print("\"\"\"\n");
 }
 
-bool PrivateGenerator::PrintBetaServicer(const ServiceDescriptor* service) {
+bool PrivateGenerator::PrintBetaServicer(const grpc_generator::Service* service) {
+  StringMap service_dict;
+  service_dict["Service"] = service->name();
   out->Print("\n\n");
-  out->Print("class Beta$Service$Servicer(object):\n", "Service",
-             service->name());
+  out->Print(service_dict, "class Beta$Service$Servicer(object):\n");
   {
     IndentScope raii_class_indent(out);
     out->Print(
@@ -243,16 +139,20 @@ bool PrivateGenerator::PrintBetaServicer(const ServiceDescriptor* service) {
         "generated\n"
         "only to ease transition from grpcio<0.15.0 to "
         "grpcio>=0.15.0.\"\"\"\n");
-    PrintAllComments(service);
+    StringVector service_comments = service->GetAllComments();
+    PrintAllComments(service_comments);
     for (int i = 0; i < service->method_count(); ++i) {
-      const MethodDescriptor* method = service->method(i);
+      auto method = service->method(i);
       grpc::string arg_name =
-          method->client_streaming() ? "request_iterator" : "request";
-      out->Print("def $Method$(self, $ArgName$, context):\n", "Method",
-                 method->name(), "ArgName", arg_name);
+          method.get()->ClientStreaming() ? "request_iterator" : "request";
+      StringMap method_dict;
+      method_dict["Method"] = method.get()->name();
+      method_dict["ArgName"] = arg_name;
+      out->Print(method_dict, "def $Method$(self, $ArgName$, context):\n");
       {
         IndentScope raii_method_indent(out);
-        PrintAllComments(method);
+        StringVector method_comments = method.get()->GetAllComments();
+        PrintAllComments(method_comments);
         out->Print("context.code(beta_interfaces.StatusCode.UNIMPLEMENTED)\n");
       }
     }
@@ -260,9 +160,11 @@ bool PrivateGenerator::PrintBetaServicer(const ServiceDescriptor* service) {
   return true;
 }
 
-bool PrivateGenerator::PrintBetaStub(const ServiceDescriptor* service) {
+bool PrivateGenerator::PrintBetaStub(const grpc_generator::Service* service) {
+  StringMap service_dict;
+  service_dict["Service"] = service->name();
   out->Print("\n\n");
-  out->Print("class Beta$Service$Stub(object):\n", "Service", service->name());
+  out->Print(service_dict, "class Beta$Service$Stub(object):\n");
   {
     IndentScope raii_class_indent(out);
     out->Print(
@@ -272,23 +174,25 @@ bool PrivateGenerator::PrintBetaStub(const ServiceDescriptor* service) {
         "generated\n"
         "only to ease transition from grpcio<0.15.0 to "
         "grpcio>=0.15.0.\"\"\"\n");
-    PrintAllComments(service);
+    StringVector service_comments = service->GetAllComments();
+    PrintAllComments(service_comments);
     for (int i = 0; i < service->method_count(); ++i) {
-      const MethodDescriptor* method = service->method(i);
+      auto method = service->method(i);
       grpc::string arg_name =
-          method->client_streaming() ? "request_iterator" : "request";
+          method.get()->ClientStreaming() ? "request_iterator" : "request";
       StringMap method_dict;
-      method_dict["Method"] = method->name();
+      method_dict["Method"] = method.get()->name();
       method_dict["ArgName"] = arg_name;
       out->Print(method_dict,
                  "def $Method$(self, $ArgName$, timeout, metadata=None, "
                  "with_call=False, protocol_options=None):\n");
       {
         IndentScope raii_method_indent(out);
-        PrintAllComments(method);
+        StringVector method_comments = method.get()->GetAllComments();
+        PrintAllComments(method_comments);
         out->Print("raise NotImplementedError()\n");
       }
-      if (!method->server_streaming()) {
+      if (!method.get()->ServerStreaming()) {
         out->Print(method_dict, "$Method$.future = None\n");
       }
     }
@@ -298,12 +202,13 @@ bool PrivateGenerator::PrintBetaStub(const ServiceDescriptor* service) {
 
 bool PrivateGenerator::PrintBetaServerFactory(
     const grpc::string& package_qualified_service_name,
-    const ServiceDescriptor* service) {
+    const grpc_generator::Service* service) {
+  StringMap service_dict;
+  service_dict["Service"] = service->name();
   out->Print("\n\n");
-  out->Print(
+  out->Print(service_dict,
       "def beta_create_$Service$_server(servicer, pool=None, "
-      "pool_size=None, default_timeout=None, maximum_timeout=None):\n",
-      "Service", service->name());
+      "pool_size=None, default_timeout=None, maximum_timeout=None):\n");
   {
     IndentScope raii_create_server_indent(out);
     out->Print(
@@ -316,41 +221,42 @@ bool PrivateGenerator::PrintBetaServerFactory(
     StringMap input_message_modules_and_classes;
     StringMap output_message_modules_and_classes;
     for (int i = 0; i < service->method_count(); ++i) {
-      const MethodDescriptor* method = service->method(i);
+      auto method = service->method(i);
       const grpc::string method_implementation_constructor =
-          grpc::string(method->client_streaming() ? "stream_" : "unary_") +
-          grpc::string(method->server_streaming() ? "stream_" : "unary_") +
+          grpc::string(method.get()->ClientStreaming() ? "stream_" : "unary_") +
+          grpc::string(method.get()->ServerStreaming() ? "stream_" : "unary_") +
           "inline";
       grpc::string input_message_module_and_class;
-      if (!GetModuleAndMessagePath(method->input_type(),
-                                   &input_message_module_and_class)) {
+      if (!method.get()->get_module_and_message_path_input(input_message_module_and_class,
+                                                           generator_file_name, generate_in_pb2_grpc)) {
         return false;
       }
       grpc::string output_message_module_and_class;
-      if (!GetModuleAndMessagePath(method->output_type(),
-                                   &output_message_module_and_class)) {
+      if (!method.get()->get_module_and_message_path_output(output_message_module_and_class,
+                                                            generator_file_name, generate_in_pb2_grpc)) {
         return false;
       }
       method_implementation_constructors.insert(
-          make_pair(method->name(), method_implementation_constructor));
+          make_pair(method.get()->name(), method_implementation_constructor));
       input_message_modules_and_classes.insert(
-          make_pair(method->name(), input_message_module_and_class));
+          make_pair(method.get()->name(), input_message_module_and_class));
       output_message_modules_and_classes.insert(
-          make_pair(method->name(), output_message_module_and_class));
+          make_pair(method.get()->name(), output_message_module_and_class));
     }
+    StringMap method_dict;
+    method_dict["PackageQualifiedServiceName"] = package_qualified_service_name;
     out->Print("request_deserializers = {\n");
     for (StringMap::iterator name_and_input_module_class_pair =
              input_message_modules_and_classes.begin();
          name_and_input_module_class_pair !=
          input_message_modules_and_classes.end();
          name_and_input_module_class_pair++) {
+      method_dict["MethodName"] = name_and_input_module_class_pair->first;
+      method_dict["InputTypeModuleAndClass"] = name_and_input_module_class_pair->second;
       IndentScope raii_indent(out);
-      out->Print(
+      out->Print(method_dict,
           "(\'$PackageQualifiedServiceName$\', \'$MethodName$\'): "
-          "$InputTypeModuleAndClass$.FromString,\n",
-          "PackageQualifiedServiceName", package_qualified_service_name,
-          "MethodName", name_and_input_module_class_pair->first,
-          "InputTypeModuleAndClass", name_and_input_module_class_pair->second);
+          "$InputTypeModuleAndClass$.FromString,\n");
     }
     out->Print("}\n");
     out->Print("response_serializers = {\n");
@@ -359,14 +265,12 @@ bool PrivateGenerator::PrintBetaServerFactory(
          name_and_output_module_class_pair !=
          output_message_modules_and_classes.end();
          name_and_output_module_class_pair++) {
+      method_dict["MethodName"] = name_and_output_module_class_pair->first;
+      method_dict["OutputTypeModuleAndClass"] = name_and_output_module_class_pair->second;
       IndentScope raii_indent(out);
-      out->Print(
+      out->Print(method_dict,
           "(\'$PackageQualifiedServiceName$\', \'$MethodName$\'): "
-          "$OutputTypeModuleAndClass$.SerializeToString,\n",
-          "PackageQualifiedServiceName", package_qualified_service_name,
-          "MethodName", name_and_output_module_class_pair->first,
-          "OutputTypeModuleAndClass",
-          name_and_output_module_class_pair->second);
+          "$OutputTypeModuleAndClass$.SerializeToString,\n");
     }
     out->Print("}\n");
     out->Print("method_implementations = {\n");
@@ -375,15 +279,14 @@ bool PrivateGenerator::PrintBetaServerFactory(
          name_and_implementation_constructor !=
          method_implementation_constructors.end();
          name_and_implementation_constructor++) {
+      method_dict["Method"] = name_and_implementation_constructor->first;
+      method_dict["Constructor"] = name_and_implementation_constructor->second;
       IndentScope raii_descriptions_indent(out);
       const grpc::string method_name =
           name_and_implementation_constructor->first;
-      out->Print(
+      out->Print(method_dict,
           "(\'$PackageQualifiedServiceName$\', \'$Method$\'): "
-          "face_utilities.$Constructor$(servicer.$Method$),\n",
-          "PackageQualifiedServiceName", package_qualified_service_name,
-          "Method", name_and_implementation_constructor->first, "Constructor",
-          name_and_implementation_constructor->second);
+          "face_utilities.$Constructor$(servicer.$Method$),\n");
     }
     out->Print("}\n");
     out->Print(
@@ -402,7 +305,7 @@ bool PrivateGenerator::PrintBetaServerFactory(
 
 bool PrivateGenerator::PrintBetaStubFactory(
     const grpc::string& package_qualified_service_name,
-    const ServiceDescriptor* service) {
+    const grpc_generator::Service* service) {
   StringMap dict;
   dict["Service"] = service->name();
   out->Print("\n\n");
@@ -421,40 +324,41 @@ bool PrivateGenerator::PrintBetaStubFactory(
     StringMap input_message_modules_and_classes;
     StringMap output_message_modules_and_classes;
     for (int i = 0; i < service->method_count(); ++i) {
-      const MethodDescriptor* method = service->method(i);
+      auto method = service->method(i);
       const grpc::string method_cardinality =
-          grpc::string(method->client_streaming() ? "STREAM" : "UNARY") + "_" +
-          grpc::string(method->server_streaming() ? "STREAM" : "UNARY");
+          grpc::string(method.get()->ClientStreaming() ? "STREAM" : "UNARY") + "_" +
+          grpc::string(method.get()->ServerStreaming() ? "STREAM" : "UNARY");
       grpc::string input_message_module_and_class;
-      if (!GetModuleAndMessagePath(method->input_type(),
-                                   &input_message_module_and_class)) {
+      if (!method.get()->get_module_and_message_path_input(input_message_module_and_class,
+                                                           generator_file_name, generate_in_pb2_grpc)) {
         return false;
       }
       grpc::string output_message_module_and_class;
-      if (!GetModuleAndMessagePath(method->output_type(),
-                                   &output_message_module_and_class)) {
+      if (!method.get()->get_module_and_message_path_output(output_message_module_and_class,
+                                                            generator_file_name, generate_in_pb2_grpc)) {
         return false;
       }
       method_cardinalities.insert(
-          make_pair(method->name(), method_cardinality));
+          make_pair(method.get()->name(), method_cardinality));
       input_message_modules_and_classes.insert(
-          make_pair(method->name(), input_message_module_and_class));
+          make_pair(method.get()->name(), input_message_module_and_class));
       output_message_modules_and_classes.insert(
-          make_pair(method->name(), output_message_module_and_class));
+          make_pair(method.get()->name(), output_message_module_and_class));
     }
+    StringMap method_dict;
+    method_dict["PackageQualifiedServiceName"] = package_qualified_service_name;
     out->Print("request_serializers = {\n");
     for (StringMap::iterator name_and_input_module_class_pair =
              input_message_modules_and_classes.begin();
          name_and_input_module_class_pair !=
          input_message_modules_and_classes.end();
          name_and_input_module_class_pair++) {
+      method_dict["MethodName"] = name_and_input_module_class_pair->first;
+      method_dict["InputTypeModuleAndClass"] = name_and_input_module_class_pair->second;
       IndentScope raii_indent(out);
-      out->Print(
+      out->Print(method_dict,
           "(\'$PackageQualifiedServiceName$\', \'$MethodName$\'): "
-          "$InputTypeModuleAndClass$.SerializeToString,\n",
-          "PackageQualifiedServiceName", package_qualified_service_name,
-          "MethodName", name_and_input_module_class_pair->first,
-          "InputTypeModuleAndClass", name_and_input_module_class_pair->second);
+          "$InputTypeModuleAndClass$.SerializeToString,\n");
     }
     out->Print("}\n");
     out->Print("response_deserializers = {\n");
@@ -463,14 +367,12 @@ bool PrivateGenerator::PrintBetaStubFactory(
          name_and_output_module_class_pair !=
          output_message_modules_and_classes.end();
          name_and_output_module_class_pair++) {
+      method_dict["MethodName"] = name_and_output_module_class_pair->first;
+      method_dict["OutputTypeModuleAndClass"] = name_and_output_module_class_pair->second;
       IndentScope raii_indent(out);
-      out->Print(
+      out->Print(method_dict,
           "(\'$PackageQualifiedServiceName$\', \'$MethodName$\'): "
-          "$OutputTypeModuleAndClass$.FromString,\n",
-          "PackageQualifiedServiceName", package_qualified_service_name,
-          "MethodName", name_and_output_module_class_pair->first,
-          "OutputTypeModuleAndClass",
-          name_and_output_module_class_pair->second);
+          "$OutputTypeModuleAndClass$.FromString,\n");
     }
     out->Print("}\n");
     out->Print("cardinalities = {\n");
@@ -478,10 +380,10 @@ bool PrivateGenerator::PrintBetaStubFactory(
              method_cardinalities.begin();
          name_and_cardinality != method_cardinalities.end();
          name_and_cardinality++) {
+      method_dict["Method"] = name_and_cardinality->first;
+      method_dict["Cardinality"] = name_and_cardinality->second;
       IndentScope raii_descriptions_indent(out);
-      out->Print("\'$Method$\': cardinality.Cardinality.$Cardinality$,\n",
-                 "Method", name_and_cardinality->first, "Cardinality",
-                 name_and_cardinality->second);
+      out->Print(method_dict, "\'$Method$\': cardinality.Cardinality.$Cardinality$,\n");
     }
     out->Print("}\n");
     out->Print(
@@ -490,23 +392,25 @@ bool PrivateGenerator::PrintBetaStubFactory(
         "request_serializers=request_serializers, "
         "response_deserializers=response_deserializers, "
         "thread_pool=pool, thread_pool_size=pool_size)\n");
-    out->Print(
+    out->Print(method_dict,
         "return beta_implementations.dynamic_stub(channel, "
         "\'$PackageQualifiedServiceName$\', "
-        "cardinalities, options=stub_options)\n",
-        "PackageQualifiedServiceName", package_qualified_service_name);
+        "cardinalities, options=stub_options)\n");
   }
   return true;
 }
 
 bool PrivateGenerator::PrintStub(
     const grpc::string& package_qualified_service_name,
-    const ServiceDescriptor* service) {
+    const grpc_generator::Service* service) {
+  StringMap dict;
+  dict["Service"] = service->name();
   out->Print("\n\n");
-  out->Print("class $Service$Stub(object):\n", "Service", service->name());
+  out->Print(dict, "class $Service$Stub(object):\n");
   {
     IndentScope raii_class_indent(out);
-    PrintAllComments(service);
+    StringVector service_comments = service->GetAllComments();
+    PrintAllComments(service_comments);
     out->Print("\n");
     out->Print("def __init__(self, channel):\n");
     {
@@ -520,35 +424,35 @@ bool PrivateGenerator::PrintStub(
       }
       out->Print("\"\"\"\n");
       for (int i = 0; i < service->method_count(); ++i) {
-        const MethodDescriptor* method = service->method(i);
+        auto method = service->method(i);
         grpc::string multi_callable_constructor =
-            grpc::string(method->client_streaming() ? "stream" : "unary") +
-            "_" + grpc::string(method->server_streaming() ? "stream" : "unary");
+            grpc::string(method.get()->ClientStreaming() ? "stream" : "unary") +
+            "_" + grpc::string(method.get()->ServerStreaming() ? "stream" : "unary");
         grpc::string request_module_and_class;
-        if (!GetModuleAndMessagePath(method->input_type(),
-                                     &request_module_and_class)) {
+        if (!method.get()->get_module_and_message_path_input(request_module_and_class,
+                                                             generator_file_name, generate_in_pb2_grpc)) {
           return false;
         }
         grpc::string response_module_and_class;
-        if (!GetModuleAndMessagePath(method->output_type(),
-                                     &response_module_and_class)) {
+        if (!method.get()->get_module_and_message_path_output(response_module_and_class,
+                                                              generator_file_name, generate_in_pb2_grpc)) {
           return false;
         }
-        out->Print("self.$Method$ = channel.$MultiCallableConstructor$(\n",
-                   "Method", method->name(), "MultiCallableConstructor",
-                   multi_callable_constructor);
+        StringMap method_dict;
+        method_dict["Method"] = method.get()->name();
+        method_dict["MultiCallableConstructor"] = multi_callable_constructor;
+        out->Print(method_dict, "self.$Method$ = channel.$MultiCallableConstructor$(\n");
         {
+          method_dict["PackageQualifiedService"] = package_qualified_service_name;
+          method_dict["RequestModuleAndClass"] = request_module_and_class;
+          method_dict["ResponseModuleAndClass"] = response_module_and_class;
           IndentScope raii_first_attribute_indent(out);
           IndentScope raii_second_attribute_indent(out);
-          out->Print("'/$PackageQualifiedService$/$Method$',\n",
-                     "PackageQualifiedService", package_qualified_service_name,
-                     "Method", method->name());
-          out->Print(
-              "request_serializer=$RequestModuleAndClass$.SerializeToString,\n",
-              "RequestModuleAndClass", request_module_and_class);
-          out->Print(
-              "response_deserializer=$ResponseModuleAndClass$.FromString,\n",
-              "ResponseModuleAndClass", response_module_and_class);
+          out->Print(method_dict, "'/$PackageQualifiedService$/$Method$',\n");
+          out->Print(method_dict,
+              "request_serializer=$RequestModuleAndClass$.SerializeToString,\n");
+          out->Print(method_dict,
+              "response_deserializer=$ResponseModuleAndClass$.FromString,\n");
           out->Print(")\n");
         }
       }
@@ -557,22 +461,28 @@ bool PrivateGenerator::PrintStub(
   return true;
 }
 
-bool PrivateGenerator::PrintServicer(const ServiceDescriptor* service) {
+bool PrivateGenerator::PrintServicer(const grpc_generator::Service* service) {
+  StringMap service_dict;
+  service_dict["Service"] = service->name();
   out->Print("\n\n");
-  out->Print("class $Service$Servicer(object):\n", "Service", service->name());
+  out->Print(service_dict, "class $Service$Servicer(object):\n");
   {
     IndentScope raii_class_indent(out);
-    PrintAllComments(service);
+    StringVector service_comments = service->GetAllComments();
+    PrintAllComments(service_comments);
     for (int i = 0; i < service->method_count(); ++i) {
-      const MethodDescriptor* method = service->method(i);
+      auto method = service->method(i);
       grpc::string arg_name =
-          method->client_streaming() ? "request_iterator" : "request";
+          method.get()->ClientStreaming() ? "request_iterator" : "request";
+      StringMap method_dict;
+      method_dict["Method"] = method.get()->name();
+      method_dict["ArgName"] = arg_name;
       out->Print("\n");
-      out->Print("def $Method$(self, $ArgName$, context):\n", "Method",
-                 method->name(), "ArgName", arg_name);
+      out->Print(method_dict, "def $Method$(self, $ArgName$, context):\n");
       {
         IndentScope raii_method_indent(out);
-        PrintAllComments(method);
+        StringVector method_comments = method.get()->GetAllComments();
+        PrintAllComments(method_comments);
         out->Print("context.set_code(grpc.StatusCode.UNIMPLEMENTED)\n");
         out->Print("context.set_details('Method not implemented!')\n");
         out->Print("raise NotImplementedError('Method not implemented!')\n");
@@ -584,10 +494,11 @@ bool PrivateGenerator::PrintServicer(const ServiceDescriptor* service) {
 
 bool PrivateGenerator::PrintAddServicerToServer(
     const grpc::string& package_qualified_service_name,
-    const ServiceDescriptor* service) {
+    const grpc_generator::Service* service) {
+  StringMap service_dict;
+  service_dict["Service"] = service->name();
   out->Print("\n\n");
-  out->Print("def add_$Service$Servicer_to_server(servicer, server):\n",
-             "Service", service->name());
+  out->Print(service_dict, "def add_$Service$Servicer_to_server(servicer, server):\n");
   {
     IndentScope raii_class_indent(out);
     out->Print("rpc_method_handlers = {\n");
@@ -595,47 +506,49 @@ bool PrivateGenerator::PrintAddServicerToServer(
       IndentScope raii_dict_first_indent(out);
       IndentScope raii_dict_second_indent(out);
       for (int i = 0; i < service->method_count(); ++i) {
-        const MethodDescriptor* method = service->method(i);
+        auto method = service->method(i);
         grpc::string method_handler_constructor =
-            grpc::string(method->client_streaming() ? "stream" : "unary") +
+            grpc::string(method.get()->ClientStreaming() ? "stream" : "unary") +
             "_" +
-            grpc::string(method->server_streaming() ? "stream" : "unary") +
+            grpc::string(method.get()->ServerStreaming() ? "stream" : "unary") +
             "_rpc_method_handler";
         grpc::string request_module_and_class;
-        if (!GetModuleAndMessagePath(method->input_type(),
-                                     &request_module_and_class)) {
+        if (!method.get()->get_module_and_message_path_input(request_module_and_class, 
+                                                             generator_file_name, generate_in_pb2_grpc)) {
           return false;
         }
         grpc::string response_module_and_class;
-        if (!GetModuleAndMessagePath(method->output_type(),
-                                     &response_module_and_class)) {
+        if (!method.get()->get_module_and_message_path_output(response_module_and_class,
+                                                              generator_file_name, generate_in_pb2_grpc)) {
           return false;
         }
-        out->Print("'$Method$': grpc.$MethodHandlerConstructor$(\n", "Method",
-                   method->name(), "MethodHandlerConstructor",
-                   method_handler_constructor);
+        StringMap method_dict;
+        method_dict["Method"] = method.get()->name();
+        method_dict["MethodHandlerConstructor"] = method_handler_constructor;
+        method_dict["RequestModuleAndClass"] = request_module_and_class;
+        method_dict["ResponseModuleAndClass"] = response_module_and_class;
+        out->Print(method_dict, "'$Method$': grpc.$MethodHandlerConstructor$(\n");
         {
           IndentScope raii_call_first_indent(out);
           IndentScope raii_call_second_indent(out);
-          out->Print("servicer.$Method$,\n", "Method", method->name());
-          out->Print(
-              "request_deserializer=$RequestModuleAndClass$.FromString,\n",
-              "RequestModuleAndClass", request_module_and_class);
-          out->Print(
+          out->Print(method_dict, "servicer.$Method$,\n");
+          out->Print(method_dict,
+              "request_deserializer=$RequestModuleAndClass$.FromString,\n");
+          out->Print(method_dict,
               "response_serializer=$ResponseModuleAndClass$.SerializeToString,"
-              "\n",
-              "ResponseModuleAndClass", response_module_and_class);
+              "\n");
         }
         out->Print("),\n");
       }
     }
+    StringMap method_dict;
+    method_dict["PackageQualifiedServiceName"] = package_qualified_service_name;
     out->Print("}\n");
     out->Print("generic_handler = grpc.method_handlers_generic_handler(\n");
     {
       IndentScope raii_call_first_indent(out);
       IndentScope raii_call_second_indent(out);
-      out->Print("'$PackageQualifiedServiceName$', rpc_method_handlers)\n",
-                 "PackageQualifiedServiceName", package_qualified_service_name);
+      out->Print(method_dict, "'$PackageQualifiedServiceName$', rpc_method_handlers)\n");
     }
     out->Print("server.add_generic_rpc_handlers((generic_handler,))\n");
   }
@@ -643,10 +556,10 @@ bool PrivateGenerator::PrintAddServicerToServer(
 }
 
 bool PrivateGenerator::PrintBetaPreamble() {
-  out->Print("from $Package$ import implementations as beta_implementations\n",
-             "Package", config.beta_package_root);
-  out->Print("from $Package$ import interfaces as beta_interfaces\n", "Package",
-             config.beta_package_root);
+  StringMap var;
+  var["Package"] = config.beta_package_root;
+  out->Print(var, "from $Package$ import implementations as beta_implementations\n");
+  out->Print(var, "from $Package$ import interfaces as beta_interfaces\n");
   out->Print("from grpc.framework.common import cardinality\n");
   out->Print(
       "from grpc.framework.interfaces.face import utilities as "
@@ -655,14 +568,16 @@ bool PrivateGenerator::PrintBetaPreamble() {
 }
 
 bool PrivateGenerator::PrintPreamble() {
-  out->Print("import $Package$\n", "Package", config.grpc_package_root);
+  StringMap var;
+  var["Package"] = config.grpc_package_root;
+  out->Print(var, "import $Package$\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);
+      auto service = file->service(i).get();
       for (int j = 0; j < service->method_count(); ++j) {
-        const MethodDescriptor* method = service->method(j);
+        const MethodDescriptor* method = (MethodDescriptor*)service->method(j).get();
         const Descriptor* types[2] = {method->input_type(),
                                       method->output_type()};
         for (int k = 0; k < 2; ++k) {
@@ -676,10 +591,12 @@ bool PrivateGenerator::PrintPreamble() {
         }
       }
     }
+    StringMap var;
     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));
+      var["ModuleName"] = std::get<0>(*it);
+      var["ModuleAlias"] = std::get<1>(*it);
+      out->Print(var, "import $ModuleName$ as $ModuleAlias$\n");
     }
   }
   return true;
@@ -691,7 +608,7 @@ bool PrivateGenerator::PrintGAServices() {
     package = package.append(".");
   }
   for (int i = 0; i < file->service_count(); ++i) {
-    const ServiceDescriptor* service = file->service(i);
+    auto service = file->service(i).get();
     grpc::string package_qualified_service_name = package + service->name();
     if (!(PrintStub(package_qualified_service_name, service) &&
           PrintServicer(service) &&
@@ -708,7 +625,7 @@ bool PrivateGenerator::PrintBetaServices() {
     package = package.append(".");
   }
   for (int i = 0; i < file->service_count(); ++i) {
-    const ServiceDescriptor* service = file->service(i);
+    auto service = file->service(i).get();
     grpc::string package_qualified_service_name = package + service->name();
     if (!(PrintBetaServicer(service) && PrintBetaStub(service) &&
           PrintBetaServerFactory(package_qualified_service_name, service) &&
@@ -718,14 +635,12 @@ bool PrivateGenerator::PrintBetaServices() {
   }
   return true;
 }
-
+  
 pair<bool, grpc::string> PrivateGenerator::GetGrpcServices() {
   grpc::string output;
   {
     // Scope the output stream so it closes and finalizes output to the string.
-    StringOutputStream output_stream(&output);
-    Printer out_printer(&output_stream, '$');
-    out = &out_printer;
+    auto out = file->CreatePrinter(&output).get();
 
     if (generate_in_pb2_grpc) {
       out->Print(
@@ -823,6 +738,7 @@ bool PythonGrpcGenerator::Generate(const FileDescriptor* file,
     *error = "Invalid proto file name. Proto file must end with .proto";
     return false;
   }
+  generator_file_name = file->name();
 
   PrivateGenerator generator(config_, file);
   if (parameter == "grpc_2_0") {