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

   1 // Copyright 2016 gRPC authors.
   2 //
   3 // Licensed under the Apache License, Version 2.0 (the "License");
   4 // you may not use this file except in compliance with the License.
   5 // You may obtain a copy of the License at
   6 //
   7 //     http://www.apache.org/licenses/LICENSE-2.0
   8 //
   9 // Unless required by applicable law or agreed to in writing, software
  10 // distributed under the License is distributed on an "AS IS" BASIS,
  11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 // See the License for the specific language governing permissions and
  13 // limitations under the License.
  14
  15 #include <google/protobuf/compiler/command_line_interface.h>
  16 #include <google/protobuf/compiler/python/python_generator.h>
  17
  18 #include "src/compiler/python_generator.h"
  19
  20 #include "grpc_tools/main.h"
  21
  22 #include <google/protobuf/compiler/code_generator.h>
  23 #include <google/protobuf/compiler/importer.h>
  24 #include <google/protobuf/descriptor.h>
  25 #include <google/protobuf/io/zero_copy_stream_impl_lite.h>
  26
  27 #include <algorithm>
  28 #include <map>
  29 #include <string>
  30 #include <tuple>
  31 #include <unordered_set>
  32 #include <vector>
  33
  34 using ::google::protobuf::FileDescriptor;
  35 using ::google::protobuf::compiler::CodeGenerator;
  36 using ::google::protobuf::compiler::DiskSourceTree;
  37 using ::google::protobuf::compiler::GeneratorContext;
  38 using ::google::protobuf::compiler::Importer;
  39 using ::google::protobuf::compiler::MultiFileErrorCollector;
  40 using ::google::protobuf::io::StringOutputStream;
  41 using ::google::protobuf::io::ZeroCopyOutputStream;
  42
  43 namespace grpc_tools {
  44 int protoc_main(int argc, char* argv[]) {
  45   google::protobuf::compiler::CommandLineInterface cli;
  46   cli.AllowPlugins("protoc-");
  47
  48   // Proto2 Python
  49   google::protobuf::compiler::python::Generator py_generator;
  50   cli.RegisterGenerator("--python_out", &py_generator,
  51                         "Generate Python source file.");
  52
  53   // gRPC Python
  54   grpc_python_generator::GeneratorConfiguration grpc_py_config;
  55   grpc_python_generator::PythonGrpcGenerator grpc_py_generator(grpc_py_config);
  56   cli.RegisterGenerator("--grpc_python_out", &grpc_py_generator,
  57                         "Generate Python source file.");
  58
  59   return cli.Run(argc, argv);
  60 }
  61
  62 namespace internal {
  63
  64 class GeneratorContextImpl : public GeneratorContext {
  65  public:
  66   GeneratorContextImpl(
  67       const std::vector<const FileDescriptor*>& parsed_files,
  68       std::vector<std::pair<std::string, std::string>>* files_out)
  69       : files_(files_out), parsed_files_(parsed_files) {}
  70
  71   ZeroCopyOutputStream* Open(const std::string& filename) {
  72     files_->emplace_back(filename, "");
  73     return new StringOutputStream(&(files_->back().second));
  74   }
  75
  76   // NOTE(rbellevi): Equivalent to Open, since all files start out empty.
  77   ZeroCopyOutputStream* OpenForAppend(const std::string& filename) {
  78     return Open(filename);
  79   }
  80
  81   // NOTE(rbellevi): Equivalent to Open, since all files start out empty.
  82   ZeroCopyOutputStream* OpenForInsert(const std::string& filename,
  83                                       const std::string& insertion_point) {
  84     return Open(filename);
  85   }
  86
  87   void ListParsedFiles(
  88       std::vector<const ::google::protobuf::FileDescriptor*>* output) {
  89     *output = parsed_files_;
  90   }
  91
  92  private:
  93   std::vector<std::pair<std::string, std::string>>* files_;
  94   const std::vector<const FileDescriptor*>& parsed_files_;
  95 };
  96
  97 class ErrorCollectorImpl : public MultiFileErrorCollector {
  98  public:
  99   ErrorCollectorImpl(std::vector<::grpc_tools::ProtocError>* errors,
 100                      std::vector<::grpc_tools::ProtocWarning>* warnings)
 101       : errors_(errors), warnings_(warnings) {}
 102
 103   void AddError(const std::string& filename, int line, int column,
 104                 const std::string& message) {
 105     errors_->emplace_back(filename, line, column, message);
 106   }
 107
 108   void AddWarning(const std::string& filename, int line, int column,
 109                   const std::string& message) {
 110     warnings_->emplace_back(filename, line, column, message);
 111   }
 112
 113  private:
 114   std::vector<::grpc_tools::ProtocError>* errors_;
 115   std::vector<::grpc_tools::ProtocWarning>* warnings_;
 116 };
 117
 118 static void calculate_transitive_closure(
 119     const FileDescriptor* descriptor,
 120     std::vector<const FileDescriptor*>* transitive_closure,
 121     std::unordered_set<const ::google::protobuf::FileDescriptor*>* visited) {
 122   for (int i = 0; i < descriptor->dependency_count(); ++i) {
 123     const FileDescriptor* dependency = descriptor->dependency(i);
 124     if (visited->find(dependency) == visited->end()) {
 125       calculate_transitive_closure(dependency, transitive_closure, visited);
 126     }
 127   }
 128   transitive_closure->push_back(descriptor);
 129   visited->insert(descriptor);
 130 }
 131
 132 }  // end namespace internal
 133
 134 static int generate_code(
 135     CodeGenerator* code_generator, char* protobuf_path,
 136     const std::vector<std::string>* include_paths,
 137     std::vector<std::pair<std::string, std::string>>* files_out,
 138     std::vector<::grpc_tools::ProtocError>* errors,
 139     std::vector<::grpc_tools::ProtocWarning>* warnings) {
 140   std::unique_ptr<internal::ErrorCollectorImpl> error_collector(
 141       new internal::ErrorCollectorImpl(errors, warnings));
 142   std::unique_ptr<DiskSourceTree> source_tree(new DiskSourceTree());
 143   for (const auto& include_path : *include_paths) {
 144     source_tree->MapPath("", include_path);
 145   }
 146   Importer importer(source_tree.get(), error_collector.get());
 147   const FileDescriptor* parsed_file = importer.Import(protobuf_path);
 148   if (parsed_file == nullptr) {
 149     return 1;
 150   }
 151   std::vector<const FileDescriptor*> transitive_closure;
 152   std::unordered_set<const FileDescriptor*> visited;
 153   internal::calculate_transitive_closure(parsed_file, &transitive_closure,
 154                                          &visited);
 155   internal::GeneratorContextImpl generator_context(transitive_closure,
 156                                                    files_out);
 157   std::string error;
 158   for (const auto descriptor : transitive_closure) {
 159     code_generator->Generate(descriptor, "", &generator_context, &error);
 160   }
 161   return 0;
 162 }
 163
 164 int protoc_get_protos(
 165     char* protobuf_path, const std::vector<std::string>* include_paths,
 166     std::vector<std::pair<std::string, std::string>>* files_out,
 167     std::vector<::grpc_tools::ProtocError>* errors,
 168     std::vector<::grpc_tools::ProtocWarning>* warnings) {
 169   ::google::protobuf::compiler::python::Generator python_generator;
 170   return generate_code(&python_generator, protobuf_path, include_paths,
 171                        files_out, errors, warnings);
 172 }
 173
 174 int protoc_get_services(
 175     char* protobuf_path, const std::vector<std::string>* include_paths,
 176     std::vector<std::pair<std::string, std::string>>* files_out,
 177     std::vector<::grpc_tools::ProtocError>* errors,
 178     std::vector<::grpc_tools::ProtocWarning>* warnings) {
 179   grpc_python_generator::GeneratorConfiguration grpc_py_config;
 180   grpc_python_generator::PythonGrpcGenerator grpc_py_generator(grpc_py_config);
 181   return generate_code(&grpc_py_generator, protobuf_path, include_paths,
 182                        files_out, errors, warnings);
 183 }
 184 }  // end namespace grpc_tools