WIP. Output protos in memory, not to STDOUT

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

@@ -19,6 +19,7 @@ from .protoc import main
 
 # TODO: Get this thing to just give me the code via an FD.
 # TODO: Figure out what to do about STDOUT pollution.
+# TODO: Search sys.path to figure out project_root automatically?
 def import_protos(proto_path, project_root):
     proto_basename = os.path.basename(proto_path)
     proto_name, _ = os.path.splitext(proto_basename)

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

@@ -16,9 +16,14 @@ from libc cimport stdlib
 
 cdef extern from "grpc_tools/main.h":
   int protoc_main(int argc, char *argv[])
+  int protoc_in_memory(char* protobuf_path, char* include_path) except +
 
 def run_main(list args not None):
   cdef char **argv = <char **>stdlib.malloc(len(args)*sizeof(char *))
   for i in range(len(args)):
     argv[i] = args[i]
   return protoc_main(len(args), argv)
+
+def run_protoc_in_memory(bytes protobuf_path, bytes include_path):
+  import sys; sys.stdout.write("cython run_protoc_in_memory\n"); sys.stdout.flush()
+  return protoc_in_memory(protobuf_path, include_path)

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

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <google/protobuf/compiler/command_line_interface.h>
 #include <google/protobuf/compiler/command_line_interface.h>
 #include <google/protobuf/compiler/python/python_generator.h>
 
@@ -19,6 +20,17 @@
 
 #include "grpc_tools/main.h"
 
+#include <google/protobuf/compiler/code_generator.h>
+#include <google/protobuf/io/zero_copy_stream_impl_lite.h>
+#include <google/protobuf/compiler/importer.h>
+#include <google/protobuf/descriptor.h>
+
+// TODO: Clang format.
+#include <vector>
+#include <map>
+#include <string>
+#include <tuple>
+
 int protoc_main(int argc, char* argv[]) {
   google::protobuf::compiler::CommandLineInterface cli;
   cli.AllowPlugins("protoc-");
@@ -36,3 +48,92 @@ int protoc_main(int argc, char* argv[]) {
 
   return cli.Run(argc, argv);
 }
+
+// TODO: Figure out what Google best practices are for internal namespace like
+// this.
+namespace detail {
+
+// TODO: Consider deduping between this and command_line_interface.cc.
+// TODO: Separate declarations and definitions.
+class GeneratorContextImpl : public ::google::protobuf::compiler::GeneratorContext {
+public:
+  GeneratorContextImpl(const std::vector<const ::google::protobuf::FileDescriptor*>& parsed_files) :
+    parsed_files_(parsed_files){}
+
+  ::google::protobuf::io::ZeroCopyOutputStream* Open(const std::string& filename) {
+    // TODO(rbellevi): Learn not to dream impossible dreams. :(
+    auto [iter, _] = files_.emplace(filename, "");
+    return new ::google::protobuf::io::StringOutputStream(&(iter->second));
+  }
+
+  // NOTE: Equivalent to Open, since all files start out empty.
+  ::google::protobuf::io::ZeroCopyOutputStream* OpenForAppend(const std::string& filename) {
+    return Open(filename);
+  }
+
+  // NOTE: Equivalent to Open, since all files start out empty.
+  ::google::protobuf::io::ZeroCopyOutputStream* OpenForInsert(
+      const std::string& filename, const std::string& insertion_point) {
+    return Open(filename);
+  }
+
+  void ListParsedFiles(std::vector<const ::google::protobuf::FileDescriptor*>* output) {
+    *output = parsed_files_;
+  }
+
+  // TODO: Figure out a method with less copying.
+  std::map<std::string, std::string>
+  GetFiles() const {
+    return files_;
+  }
+
+private:
+  std::map<std::string, std::string> files_;
+  const std::vector<const ::google::protobuf::FileDescriptor*>& parsed_files_;
+};
+
+class ErrorCollectorImpl : public ::google::protobuf::compiler::MultiFileErrorCollector {
+ public:
+  ErrorCollectorImpl() {}
+  ~ErrorCollectorImpl() {}
+
+  // implements ErrorCollector ---------------------------------------
+  void AddError(const std::string& filename, int line, int column,
+                const std::string& message) {
+    // TODO: Implement.
+  }
+
+  void AddWarning(const std::string& filename, int line, int column,
+                  const std::string& message) {
+    // TODO: Implement.
+  }
+};
+
+} // end namespace detail
+
+#include <iostream>
+
+int protoc_in_memory(char* protobuf_path, char* include_path) {
+  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());
+  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);
+  detail::GeneratorContextImpl generator_context({parsed_file});
+  std::string error;
+  ::google::protobuf::compiler::python::Generator python_generator;
+  python_generator.Generate(parsed_file, "", &generator_context, &error);
+  for (const auto& [filename, contents] : generator_context.GetFiles()) {
+    std::cout << "# File: " << filename << std::endl;
+    std::cout << contents << std::endl;
+    std::cout << std::endl;
+  }
+  std::cout << std::flush;
+  // TODO: Come up with a better error reporting mechanism than this.
+  return 0;
+}

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

@@ -16,3 +16,4 @@
 // We declare `protoc_main` here since we want access to it from Cython as an
 // extern but *without* triggering a dllimport declspec when on Windows.
 int protoc_main(int argc, char *argv[]);
+int protoc_in_memory(char* protobuf_path, char* include_path);

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

@@ -29,6 +29,9 @@ def main(command_arguments):
     command_arguments = [argument.encode() for argument in command_arguments]
     return _protoc_compiler.run_main(command_arguments)
 
+def run_protoc_in_memory(protobuf_path, include_path):
+  return _protoc_compiler.run_protoc_in_memory(protobuf_path.encode('ascii'), include_path.encode('ascii'))
+
 
 if __name__ == '__main__':
     proto_include = pkg_resources.resource_filename('grpc_tools', '_proto')

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

@@ -4,24 +4,33 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import importlib
-
 import unittest
 import grpc_tools
 
-import os
-
 
 
 class ProtocTest(unittest.TestCase):
 
-    def test_protoc(self):
-        # TODO: Get this thing to just give me the code via an FD.
-        # TODO: Figure out what to do about STDOUT pollution.
-        # TODO: How do we convert protoc failure into a Python error?
-        protos, services = grpc_tools.import_protos("grpc_tools/simple.proto", "tools/distrib/python/grpcio_tools/")
-        print(dir(protos))
-        print(dir(services))
+    # def test_import_protos(self):
+    #     protos, services = grpc_tools.import_protos("grpc_tools/simple.proto", "tools/distrib/python/grpcio_tools/")
+    #     print(dir(protos))
+    #     print(dir(services))
+
+    # # TODO: Ensure that we don't pollute STDOUT by invoking protoc.
+    # def test_stdout_pollution(self):
+    #     pass
+
+    def test_protoc_in_memory(self):
+        print("Running test_protoc_in_memory")
+        from grpc_tools import protoc
+        import os
+        original_dir = os.getcwd()
+        # TODO: Completely get rid of this chdir stuff.
+        os.chdir(os.path.join(original_dir, "tools/distrib/python/grpcio_tools/"))
+        protoc.run_protoc_in_memory("grpc_tools/simple.proto", "")
+        os.chdir(original_dir)
+        import sys; sys.stdout.flush()
+        print("Got to end of test.")
 
 
 if __name__ == '__main__':