Merge pull request #4699 from BSBandme/add_proto2_to_proto3_plugin
Add proto2 to proto3 utilpull/1817/merge
Yilun Chong
7 years ago
committed by
GitHub
commit
a4d16ed886
7 changed files with 356 additions and 46 deletions
@ -0,0 +1,64 @@ |
||||
#ifndef PROTOBUF_BENCHMARKS_UTIL_DATA_PROTO2_TO_PROTO3_UTIL_H_ |
||||
#define PROTOBUF_BENCHMARKS_UTIL_DATA_PROTO2_TO_PROTO3_UTIL_H_ |
||||
|
||||
#include "google/protobuf/message.h" |
||||
#include "google/protobuf/descriptor.h" |
||||
|
||||
using google::protobuf::FieldDescriptor; |
||||
using google::protobuf::Message; |
||||
using google::protobuf::Reflection; |
||||
|
||||
namespace google { |
||||
namespace protobuf { |
||||
namespace util { |
||||
|
||||
class DataStripper { |
||||
public: |
||||
void StripMessage(Message *message) { |
||||
std::vector<const FieldDescriptor*> set_fields; |
||||
const Reflection* reflection = message->GetReflection(); |
||||
reflection->ListFields(*message, &set_fields); |
||||
|
||||
for (size_t i = 0; i < set_fields.size(); i++) { |
||||
const FieldDescriptor* field = set_fields[i]; |
||||
if (ShouldBeClear(field)) { |
||||
reflection->ClearField(message, field); |
||||
continue; |
||||
} |
||||
if (field->type() == FieldDescriptor::TYPE_MESSAGE) { |
||||
if (field->is_repeated()) { |
||||
for (int j = 0; j < reflection->FieldSize(*message, field); j++) { |
||||
StripMessage(reflection->MutableRepeatedMessage(message, field, j)); |
||||
} |
||||
} else { |
||||
StripMessage(reflection->MutableMessage(message, field)); |
||||
} |
||||
} |
||||
} |
||||
|
||||
reflection->MutableUnknownFields(message)->Clear(); |
||||
} |
||||
private: |
||||
virtual bool ShouldBeClear(const FieldDescriptor *field) = 0; |
||||
}; |
||||
|
||||
class GogoDataStripper : public DataStripper { |
||||
private: |
||||
virtual bool ShouldBeClear(const FieldDescriptor *field) { |
||||
return field->type() == FieldDescriptor::TYPE_GROUP; |
||||
} |
||||
}; |
||||
|
||||
class Proto3DataStripper : public DataStripper { |
||||
private: |
||||
virtual bool ShouldBeClear(const FieldDescriptor *field) { |
||||
return field->type() == FieldDescriptor::TYPE_GROUP || |
||||
field->is_extension(); |
||||
} |
||||
}; |
||||
|
||||
} // namespace util
|
||||
} // namespace protobuf
|
||||
} // namespace google
|
||||
|
||||
#endif // PROTOBUF_BENCHMARKS_UTIL_DATA_PROTO2_TO_PROTO3_UTIL_H_
|
||||
@ -0,0 +1,74 @@ |
||||
#include "benchmarks.pb.h" |
||||
#include "datasets/google_message1/proto2/benchmark_message1_proto2.pb.h" |
||||
#include "datasets/google_message1/proto3/benchmark_message1_proto3.pb.h" |
||||
#include "datasets/google_message2/benchmark_message2.pb.h" |
||||
#include "datasets/google_message3/benchmark_message3.pb.h" |
||||
#include "datasets/google_message4/benchmark_message4.pb.h" |
||||
#include "data_proto2_to_proto3_util.h" |
||||
|
||||
#include <fstream> |
||||
|
||||
using google::protobuf::util::Proto3DataStripper; |
||||
|
||||
std::string ReadFile(const std::string& name) { |
||||
std::ifstream file(name.c_str()); |
||||
GOOGLE_CHECK(file.is_open()) << "Couldn't find file '" |
||||
<< name |
||||
<< "', please make sure you are running this command from the benchmarks" |
||||
<< " directory.\n"; |
||||
return std::string((std::istreambuf_iterator<char>(file)), |
||||
std::istreambuf_iterator<char>()); |
||||
} |
||||
|
||||
int main(int argc, char *argv[]) { |
||||
if (argc % 2 == 0 || argc == 1) { |
||||
std::cerr << "Usage: [input_files] [output_file_names] where " << |
||||
"input_files are one to one mapping to output_file_names." << |
||||
std::endl; |
||||
return 1; |
||||
} |
||||
|
||||
for (int i = argc / 2; i > 0; i--) { |
||||
const std::string &input_file = argv[i]; |
||||
const std::string &output_file = argv[i + argc / 2]; |
||||
|
||||
std::cerr << "Generating " << input_file |
||||
<< " to " << output_file << std::endl; |
||||
benchmarks::BenchmarkDataset dataset; |
||||
Message* message; |
||||
std::string dataset_payload = ReadFile(input_file); |
||||
GOOGLE_CHECK(dataset.ParseFromString(dataset_payload)) |
||||
<< "Can' t parse data file " << input_file; |
||||
|
||||
if (dataset.message_name() == "benchmarks.proto3.GoogleMessage1") { |
||||
message = new benchmarks::proto3::GoogleMessage1; |
||||
} else if (dataset.message_name() == "benchmarks.proto2.GoogleMessage1") { |
||||
message = new benchmarks::proto2::GoogleMessage1; |
||||
} else if (dataset.message_name() == "benchmarks.proto2.GoogleMessage2") { |
||||
message = new benchmarks::proto2::GoogleMessage2; |
||||
} else if (dataset.message_name() == |
||||
"benchmarks.google_message3.GoogleMessage3") { |
||||
message = new benchmarks::google_message3::GoogleMessage3; |
||||
} else if (dataset.message_name() == |
||||
"benchmarks.google_message4.GoogleMessage4") { |
||||
message = new benchmarks::google_message4::GoogleMessage4; |
||||
} else { |
||||
std::cerr << "Unknown message type: " << dataset.message_name(); |
||||
exit(1); |
||||
} |
||||
|
||||
for (int i = 0; i < dataset.payload_size(); i++) { |
||||
message->ParseFromString(dataset.payload(i)); |
||||
Proto3DataStripper stripper; |
||||
stripper.StripMessage(message); |
||||
dataset.set_payload(i, message->SerializeAsString()); |
||||
} |
||||
|
||||
std::ofstream ofs(output_file); |
||||
ofs << dataset.SerializeAsString(); |
||||
ofs.close(); |
||||
} |
||||
|
||||
|
||||
return 0; |
||||
} |
||||
@ -0,0 +1,115 @@ |
||||
#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) { |
||||
return filename.substr(0, filename.rfind(".proto")); |
||||
} |
||||
|
||||
DescriptorPool* GetPool() { |
||||
static DescriptorPool *pool = new DescriptorPool(); |
||||
return pool; |
||||
} |
||||
|
||||
} // namespace
|
||||
|
||||
class Proto2ToProto3Generator final : public CodeGenerator { |
||||
public: |
||||
bool GenerateAll(const std::vector<const FileDescriptor*>& files, |
||||
const string& parameter, |
||||
GeneratorContext* context, |
||||
string* error) const { |
||||
for (int i = 0; i < files.size(); i++) { |
||||
for (auto file : files) { |
||||
if (CanGenerate(file)) { |
||||
Generate(file, parameter, context, error); |
||||
break; |
||||
} |
||||
} |
||||
} |
||||
|
||||
return true; |
||||
} |
||||
|
||||
bool Generate(const FileDescriptor* file, |
||||
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")); |
||||
string content = GetPool()->BuildFile(new_file)->DebugString(); |
||||
Printer printer(output.get(), '$'); |
||||
printer.WriteRaw(content.c_str(), content.size()); |
||||
|
||||
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; |
||||
} |
||||
} |
||||
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); |
||||
} |
||||
Loading…
Reference in new issue