protoc-gen-proto2_to_proto3.cc 3.38 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25
#include "google/protobuf/compiler/code_generator.h"
#include "google/protobuf/io/zero_copy_stream.h"
#include "google/protobuf/io/printer.h"
#include "google/protobuf/descriptor.h"
#include "google/protobuf/descriptor.pb.h"
#include "schema_proto2_to_proto3_util.h"

#include "google/protobuf/compiler/plugin.h"

using google::protobuf::FileDescriptorProto;
using google::protobuf::FileDescriptor;
using google::protobuf::DescriptorPool;
using google::protobuf::io::Printer;
using google::protobuf::util::SchemaGroupStripper;
using google::protobuf::util::EnumScrubber;
using google::protobuf::util::ExtensionStripper;
using google::protobuf::util::FieldScrubber;

namespace google {
namespace protobuf {
namespace compiler {

namespace {

string StripProto(string filename) {
Yilun Chong's avatar
Yilun Chong committed
26
  return filename.substr(0, filename.rfind(".proto"));
27 28
}

Yilun Chong's avatar
Yilun Chong committed
29 30 31 32
DescriptorPool* GetPool() {
  static DescriptorPool *pool = new DescriptorPool();
  return pool;
}
33 34 35

}  // namespace

Yilun Chong's avatar
Yilun Chong committed
36
class Proto2ToProto3Generator final : public CodeGenerator {
37
 public:
Yilun Chong's avatar
Yilun Chong committed
38
  bool GenerateAll(const std::vector<const FileDescriptor*>& files,
39 40 41 42 43
                           const string& parameter,
                           GeneratorContext* context,
                           string* error) const {
    for (int i = 0; i < files.size(); i++) {
      for (auto file : files) {
Yilun Chong's avatar
Yilun Chong committed
44
        if (CanGenerate(file)) {
45 46 47 48 49 50 51 52 53
          Generate(file, parameter, context, error);
          break;
        }
      }
    }

    return true;
  }

Yilun Chong's avatar
Yilun Chong committed
54
  bool Generate(const FileDescriptor* file,
55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75
                        const string& parameter,
                        GeneratorContext* context,
                        string* error) const {
    FileDescriptorProto new_file;
    file->CopyTo(&new_file);
    SchemaGroupStripper::StripFile(file, &new_file);

    EnumScrubber enum_scrubber;
    enum_scrubber.ScrubFile(&new_file);
    ExtensionStripper::StripFile(&new_file);
    FieldScrubber::ScrubFile(&new_file);
    new_file.set_syntax("proto3");

    string filename = file->name();
    string basename = StripProto(filename);

    std::vector<std::pair<string,string>> option_pairs;
    ParseGeneratorParameter(parameter, &option_pairs);

    std::unique_ptr<google::protobuf::io::ZeroCopyOutputStream> output(
        context->Open(basename + ".proto"));
Yilun Chong's avatar
Yilun Chong committed
76
    string content = GetPool()->BuildFile(new_file)->DebugString();
77 78 79
    Printer printer(output.get(), '$');
    printer.WriteRaw(content.c_str(), content.size());

Yilun Chong's avatar
Yilun Chong committed
80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103
    return true;
  }
 private:
  bool CanGenerate(const FileDescriptor* file) const {
    if (GetPool()->FindFileByName(file->name()) != nullptr) {
      return false;
    }
    for (int j = 0; j < file->dependency_count(); j++) {
      if (GetPool()->FindFileByName(file->dependency(j)->name()) == nullptr) {
        return false;
      }
    }
    for (int j = 0; j < file->public_dependency_count(); j++) {
      if (GetPool()->FindFileByName(
          file->public_dependency(j)->name()) == nullptr) {
        return false;
      }
    }
    for (int j = 0; j < file->weak_dependency_count(); j++) {
      if (GetPool()->FindFileByName(
          file->weak_dependency(j)->name()) == nullptr) {
        return false;
      }
    }
104 105 106 107 108 109 110 111 112 113 114 115
    return true;
  }
};

}  // namespace compiler
}  // namespace protobuf
}  // namespace google

int main(int argc, char* argv[]) {
  google::protobuf::compiler::Proto2ToProto3Generator generator;
  return google::protobuf::compiler::PluginMain(argc, argv, &generator);
}