From 9cc62a62088de988b13c16595abb48c04cceefd5 Mon Sep 17 00:00:00 2001 From: Richard Belleville Date: Tue, 26 Nov 2019 15:23:08 -0800 Subject: [PATCH] WIP. Output protos in memory, not to STDOUT --- .../grpcio_tools/grpc_tools/__init__.py | 1 + .../grpc_tools/_protoc_compiler.pyx | 5 + .../python/grpcio_tools/grpc_tools/main.cc | 101 ++++++++++++++++++ .../python/grpcio_tools/grpc_tools/main.h | 1 + .../python/grpcio_tools/grpc_tools/protoc.py | 3 + .../grpcio_tools/grpc_tools/protoc_test.py | 31 ++++-- 6 files changed, 131 insertions(+), 11 deletions(-) diff --git a/tools/distrib/python/grpcio_tools/grpc_tools/__init__.py b/tools/distrib/python/grpcio_tools/grpc_tools/__init__.py index d151994fc90..3db29b406d3 100644 --- a/tools/distrib/python/grpcio_tools/grpc_tools/__init__.py +++ b/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) diff --git a/tools/distrib/python/grpcio_tools/grpc_tools/_protoc_compiler.pyx b/tools/distrib/python/grpcio_tools/grpc_tools/_protoc_compiler.pyx index 481b5a4f2bf..d94d650d214 100644 --- a/tools/distrib/python/grpcio_tools/grpc_tools/_protoc_compiler.pyx +++ b/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 = 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) diff --git a/tools/distrib/python/grpcio_tools/grpc_tools/main.cc b/tools/distrib/python/grpcio_tools/grpc_tools/main.cc index 65933865029..072db21d415 100644 --- a/tools/distrib/python/grpcio_tools/grpc_tools/main.cc +++ b/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 #include #include @@ -19,6 +20,17 @@ #include "grpc_tools/main.h" +#include +#include +#include +#include + +// TODO: Clang format. +#include +#include +#include +#include + 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& 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* output) { + *output = parsed_files_; + } + + // TODO: Figure out a method with less copying. + std::map + GetFiles() const { + return files_; + } + +private: + std::map files_; + const std::vector& 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 + +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 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; +} diff --git a/tools/distrib/python/grpcio_tools/grpc_tools/main.h b/tools/distrib/python/grpcio_tools/grpc_tools/main.h index 9a1df2028bb..22e286826f6 100644 --- a/tools/distrib/python/grpcio_tools/grpc_tools/main.h +++ b/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); diff --git a/tools/distrib/python/grpcio_tools/grpc_tools/protoc.py b/tools/distrib/python/grpcio_tools/grpc_tools/protoc.py index 18a5ef22874..60fbe1bddca 100644 --- a/tools/distrib/python/grpcio_tools/grpc_tools/protoc.py +++ b/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') diff --git a/tools/distrib/python/grpcio_tools/grpc_tools/protoc_test.py b/tools/distrib/python/grpcio_tools/grpc_tools/protoc_test.py index 01050d703ab..b227c66ccfd 100644 --- a/tools/distrib/python/grpcio_tools/grpc_tools/protoc_test.py +++ b/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__':