Clean up service import

tools/distrib/python/grpcio_tools/grpc_tools/_protoc_compiler.pyx

@@ -17,6 +17,8 @@ from libcpp.map cimport map
 from libcpp.vector cimport vector
 from libcpp.string cimport string
 
+from cython.operator cimport dereference
+
 import warnings
 
 cdef extern from "grpc_tools/main.h":
@@ -74,22 +76,24 @@ cdef _c_protoc_error_to_protoc_error(cProtocError c_protoc_error):
 cdef _c_protoc_warning_to_protoc_warning(cProtocWarning c_protoc_warning):
     return ProtocWarning(c_protoc_warning.filename, c_protoc_warning.line, c_protoc_warning.column, c_protoc_warning.message)
 
-def get_protos(bytes protobuf_path, bytes include_path):
-  cdef map[string, string] files
-  cdef vector[cProtocError] errors
-  # NOTE: Abbreviated name used to shadowing of the module name.
-  cdef vector[cProtocWarning] wrnings
-  return_value = protoc_get_protos(protobuf_path, include_path, &files, &errors, &wrnings)
-  for warning in wrnings:
+cdef _handle_errors(int rc, vector[cProtocError]* errors, vector[cProtocWarning]* wrnings, bytes protobuf_path):
+  for warning in dereference(wrnings):
       warnings.warn(_c_protoc_warning_to_protoc_warning(warning))
-  if return_value != 0:
-    if errors.size() != 0:
-       py_errors = [_c_protoc_error_to_protoc_error(c_error) for c_error in errors]
+  if rc != 0:
+    if dereference(errors).size() != 0:
+       py_errors = [_c_protoc_error_to_protoc_error(c_error) for c_error in dereference(errors)]
        # TODO: Come up with a good system for printing multiple errors from
        # protoc.
        raise Exception(py_errors)
     raise Exception("An unknown error occurred while compiling {}".format(protobuf_path))
 
+def get_protos(bytes protobuf_path, bytes include_path):
+  cdef map[string, string] files
+  cdef vector[cProtocError] errors
+  # NOTE: Abbreviated name used to shadowing of the module name.
+  cdef vector[cProtocWarning] wrnings
+  rc = protoc_get_protos(protobuf_path, include_path, &files, &errors, &wrnings)
+  _handle_errors(rc, &errors, &wrnings, protobuf_path)
   return files
 
 def get_services(bytes protobuf_path, bytes include_path):
@@ -97,16 +101,7 @@ def get_services(bytes protobuf_path, bytes include_path):
   cdef vector[cProtocError] errors
   # NOTE: Abbreviated name used to shadowing of the module name.
   cdef vector[cProtocWarning] wrnings
-  return_value = protoc_get_services(protobuf_path, include_path, &files, &errors, &wrnings)
-  for warning in wrnings:
-      warnings.warn(_c_protoc_warning_to_protoc_warning(warning))
-  if return_value != 0:
-    if errors.size() != 0:
-       py_errors = [_c_protoc_error_to_protoc_error(c_error) for c_error in errors]
-       # TODO: Come up with a good system for printing multiple errors from
-       # protoc.
-       raise Exception(py_errors)
-    raise Exception("An unknown error occurred while compiling {}".format(protobuf_path))
-
+  rc = protoc_get_services(protobuf_path, include_path, &files, &errors, &wrnings)
+  _handle_errors(rc, &errors, &wrnings, protobuf_path)
   return files
 

tools/distrib/python/grpcio_tools/grpc_tools/main.cc

@@ -114,14 +114,13 @@ private:
 
 } // end namespace detail
 
-int protoc_get_protos(char* protobuf_path,
-                     char* include_path,
-                     std::map<std::string, std::string>* files_out,
-                     std::vector<ProtocError>* errors,
-                     std::vector<ProtocWarning>* warnings)
+static int generate_code(::google::protobuf::compiler::CodeGenerator* code_generator,
+                         char* protobuf_path,
+                         char* include_path,
+                         std::map<std::string, std::string>* files_out,
+                         std::vector<ProtocError>* errors,
+                         std::vector<ProtocWarning>* warnings)
 {
-  std::cout << "C++ protoc_in_memory" << std::endl << std::flush;
-  // TODO: Create parsed_files.
   std::string protobuf_filename(protobuf_path);
   std::unique_ptr<detail::ErrorCollectorImpl> error_collector(new detail::ErrorCollectorImpl(errors, warnings));
   std::unique_ptr<::google::protobuf::compiler::DiskSourceTree> source_tree(new ::google::protobuf::compiler::DiskSourceTree());
@@ -137,33 +136,27 @@ int protoc_get_protos(char* protobuf_path,
   std::string error;
   ::google::protobuf::compiler::python::Generator python_generator;
   python_generator.Generate(parsed_file, "", &generator_context, &error);
-  // TODO: Come up with a better error reporting mechanism than this.
   return 0;
 }
 
+int protoc_get_protos(char* protobuf_path,
+                     char* include_path,
+                     std::map<std::string, std::string>* files_out,
+                     std::vector<ProtocError>* errors,
+                     std::vector<ProtocWarning>* warnings)
+{
+  ::google::protobuf::compiler::python::Generator python_generator;
+  return generate_code(&python_generator, protobuf_path, include_path, files_out, errors, warnings);
+}
+
 int protoc_get_services(char* protobuf_path,
                      char* include_path,
                      std::map<std::string, std::string>* files_out,
                      std::vector<ProtocError>* errors,
                      std::vector<ProtocWarning>* warnings)
 {
-  // TODO: Create parsed_files.
-  std::string protobuf_filename(protobuf_path);
-  std::unique_ptr<detail::ErrorCollectorImpl> error_collector(new detail::ErrorCollectorImpl(errors, warnings));
-  std::unique_ptr<::google::protobuf::compiler::DiskSourceTree> source_tree(new ::google::protobuf::compiler::DiskSourceTree());
-  // NOTE: This is equivalent to "--proto_path=."
-  source_tree->MapPath("", ".");
-  // TODO: Figure out more advanced virtual path mapping.
-  ::google::protobuf::compiler::Importer importer(source_tree.get(), error_collector.get());
-  const ::google::protobuf::FileDescriptor* parsed_file = importer.Import(protobuf_filename);
-  if (parsed_file == nullptr) {
-    return 1;
-  }
-  detail::GeneratorContextImpl generator_context({parsed_file}, files_out);
-  std::string error;
   grpc_python_generator::GeneratorConfiguration grpc_py_config;
   grpc_python_generator::PythonGrpcGenerator grpc_py_generator(grpc_py_config);
-  grpc_py_generator.Generate(parsed_file, "", &generator_context, &error);
-  // TODO: Come up with a better error reporting mechanism than this.
+  return generate_code(&grpc_py_generator, protobuf_path, include_path, files_out, errors, warnings);
   return 0;
 }

tools/distrib/python/grpcio_tools/grpc_tools/protoc.py

@@ -34,8 +34,7 @@ def main(command_arguments):
     command_arguments = [argument.encode() for argument in command_arguments]
     return _protoc_compiler.run_main(command_arguments)
 
-def get_protos(protobuf_path, include_path):
-  files = _protoc_compiler.get_protos(protobuf_path.encode('ascii'), include_path.encode('ascii'))
+def _import_modules_from_files(files):
   modules = []
   # TODO: Ensure pointer equality between two invocations of this function.
   for filename, code in six.iteritems(files):
@@ -51,20 +50,13 @@ def get_protos(protobuf_path, include_path):
     sys.modules[module_name] = module
   return tuple(modules)
 
+def get_protos(protobuf_path, include_path):
+  files = _protoc_compiler.get_protos(protobuf_path.encode('ascii'), include_path.encode('ascii'))
+  return _import_modules_from_files(files)
+
 def get_services(protobuf_path, include_path):
   files = _protoc_compiler.get_services(protobuf_path.encode('ascii'), include_path.encode('ascii'))
-  modules = []
-  # TODO: Ensure pointer equality between two invocations of this function.
-  for filename, code in six.iteritems(files):
-    base_name = os.path.basename(filename.decode('ascii'))
-    proto_name, _ = os.path.splitext(base_name)
-    anchor_package = ".".join(os.path.normpath(os.path.dirname(filename.decode('ascii'))).split(os.sep))
-    module_name = "{}.{}".format(anchor_package, proto_name)
-    module = imp.new_module(module_name)
-    six.exec_(code, module.__dict__)
-    modules.append(module)
-    sys.modules[module_name] = module
-  return tuple(modules)
+  return _import_modules_from_files(files)
 
 
 if __name__ == '__main__':