From 796d00c003112be8e09e7f2a11035ff290ff5054 Mon Sep 17 00:00:00 2001 From: iamqizhao Date: Thu, 29 Jan 2015 18:07:54 -0800 Subject: [PATCH] Make grpc-go proto plugin better: i) support external imports; ii) do not generate *_grpc.pb.go file if there is no service defined; ii) add more comment to gen.sh --- rpc/compiler/gen.sh | 12 +++-- rpc/compiler/go_generator.cc | 96 ++++++++++++++++++++++++++++++------ rpc/compiler/go_plugin.cc | 4 ++ 3 files changed, 95 insertions(+), 17 deletions(-) diff --git a/rpc/compiler/gen.sh b/rpc/compiler/gen.sh index 13fc58b4..cb56094c 100755 --- a/rpc/compiler/gen.sh +++ b/rpc/compiler/gen.sh @@ -1,8 +1,14 @@ #!/bin/bash -# Please run this under the same directory where the input proto stays. The -# output will stay in the same directory. If this is not the behavior you want, -# feel free to make your own scripts. +# This script serves as an example to demonstrate how to generate the gRPC-Go +# interface and the related messages. +# +# We suggest the importing paths in proto file are relative to $GOPATH/src and +# this script should be run at $GOPATH/src. +# +# If this is not what you need, feel free to make your own scripts. Again, this +# script is for demonstration purpose. +# locProtocGenGo=$1 locGoPlugIn=$2 proto=$3 diff --git a/rpc/compiler/go_generator.cc b/rpc/compiler/go_generator.cc index e64c658c..21484f3b 100644 --- a/rpc/compiler/go_generator.cc +++ b/rpc/compiler/go_generator.cc @@ -79,12 +79,34 @@ string LowerCaseService(const string& service) { return ret; } +std::string BadToUnderscore(std::string str) { + for (unsigned i = 0; i < str.size(); ++i) { + if (!std::isalnum(str[i])) { + str[i] = '_'; + } + } + return str; +} + +const string GetFullName(const string& selfPkg, + const string& msgPkg, + const string& msgName) { + if (selfPkg == msgPkg) { + return msgName; + } + return BadToUnderscore(msgPkg) + "." + msgName; +} + void PrintClientMethodDef(google::protobuf::io::Printer* printer, const google::protobuf::MethodDescriptor* method, map* vars) { (*vars)["Method"] = method->name(); - (*vars)["Request"] = method->input_type()->name(); - (*vars)["Response"] = method->output_type()->name(); + (*vars)["Request"] = GetFullName((*vars)["PackageName"], + method->input_type()->file()->package(), + method->input_type()->name()); + (*vars)["Response"] = GetFullName((*vars)["PackageName"], + method->output_type()->file()->package(), + method->output_type()->name()); if (NoStreaming(method)) { printer->Print(*vars, "\t$Method$(ctx context.Context, in *$Request$, opts " @@ -110,9 +132,12 @@ void PrintClientMethodImpl(google::protobuf::io::Printer* printer, const google::protobuf::MethodDescriptor* method, map* vars) { (*vars)["Method"] = method->name(); - (*vars)["Request"] = method->input_type()->name(); - (*vars)["Response"] = method->output_type()->name(); - + (*vars)["Request"] = GetFullName((*vars)["PackageName"], + method->input_type()->file()->package(), + method->input_type()->name()); + (*vars)["Response"] = GetFullName((*vars)["PackageName"], + method->output_type()->file()->package(), + method->output_type()->name()); if (NoStreaming(method)) { printer->Print( *vars, @@ -281,8 +306,12 @@ void PrintServerMethodDef(google::protobuf::io::Printer* printer, const google::protobuf::MethodDescriptor* method, map* vars) { (*vars)["Method"] = method->name(); - (*vars)["Request"] = method->input_type()->name(); - (*vars)["Response"] = method->output_type()->name(); + (*vars)["Request"] = GetFullName((*vars)["PackageName"], + method->input_type()->file()->package(), + method->input_type()->name()); + (*vars)["Response"] = GetFullName((*vars)["PackageName"], + method->output_type()->file()->package(), + method->output_type()->name()); if (NoStreaming(method)) { printer->Print( *vars, @@ -301,8 +330,12 @@ void PrintServerHandler(google::protobuf::io::Printer* printer, const google::protobuf::MethodDescriptor* method, map* vars) { (*vars)["Method"] = method->name(); - (*vars)["Request"] = method->input_type()->name(); - (*vars)["Response"] = method->output_type()->name(); + (*vars)["Request"] = GetFullName((*vars)["PackageName"], + method->input_type()->file()->package(), + method->input_type()->name()); + (*vars)["Response"] = GetFullName((*vars)["PackageName"], + method->output_type()->file()->package(), + method->output_type()->name()); if (NoStreaming(method)) { printer->Print( *vars, @@ -480,13 +513,46 @@ void PrintServer(google::protobuf::io::Printer* printer, "}\n\n"); } -std::string BadToUnderscore(std::string str) { - for (unsigned i = 0; i < str.size(); ++i) { - if (!std::isalnum(str[i])) { - str[i] = '_'; +void PrintMessageImports( + google::protobuf::io::Printer* printer, + const google::protobuf::FileDescriptor* file, + map* vars) { + set descs; + set importedPkgs; + for (int i = 0; i < file->service_count(); ++i) { + const google::protobuf::ServiceDescriptor* service = file->service(i); + for (int j = 0; j < service->method_count(); ++j) { + const google::protobuf::MethodDescriptor* method = service->method(i); + // Remove duplicated imports. + if (importedPkgs.find( + method->input_type()->file()->package()) == importedPkgs.end()) { + descs.insert(method->input_type()->file()); + importedPkgs.insert(method->input_type()->file()->package()); + } + if (importedPkgs.find( + method->output_type()->file()->package()) == importedPkgs.end()) { + descs.insert(method->output_type()->file()); + importedPkgs.insert(method->output_type()->file()->package()); + } } } - return str; + + for (auto fd : descs) { + if (fd->package() == (*vars)["PackageName"]) { + continue; + } + string name = fd->name(); + string import_path = "import \""; + if (name.find('/') == string::npos) { + // Assume all the proto in the same directory belong to the same package. + continue; + } else { + import_path += name.substr(0, name.find_last_of('/')) + "\""; + } + printer->Print(import_path.c_str()); + printer->Print("\n"); + } + printer->Print("\n"); } string GetServices(const google::protobuf::FileDescriptor* file) { @@ -512,6 +578,8 @@ string GetServices(const google::protobuf::FileDescriptor* file) { "\tproto \"github.com/golang/protobuf/proto\"\n" ")\n\n"); + PrintMessageImports(&printer, file, &vars); + // $Package$ is used to fully qualify method names. vars["Package"] = file->package(); if (!file->package().empty()) { diff --git a/rpc/compiler/go_plugin.cc b/rpc/compiler/go_plugin.cc index bcd637d2..79423164 100644 --- a/rpc/compiler/go_plugin.cc +++ b/rpc/compiler/go_plugin.cc @@ -58,6 +58,10 @@ class GoGrpcGenerator : public google::protobuf::compiler::CodeGenerator { const string& parameter, google::protobuf::compiler::GeneratorContext* context, string* error) const { + if (file->service_count() <= 0) { + // Do not generate anything if there is no rpc service defined. + return true; + } // Get output file name. string file_name; if (file->name().size() > 6 &&