4.2.2 定制代码生成插件
Protobuf的protoc编译器是通过插件机制实现对不同语言的支持。比如protoc命令出现--xxx_out格式的参数，那么protoc将首先查询是否有内置的xxx插件，如果没有内置的xxx插件那么将继续查询当前系统中是否存在protoc-gen-xxx命名的可执行程序，最终通过查询到的插件生成代码。对于Go语言的protoc-gen-go插件来说，里面又实现了一层静态插件系统。比如protoc-gen-go内置了一个gRPC插件，用户可以通过--go_out=plugins=grpc参数来生成gRPC相关代码，否则只会针对message生成相关代码。
参考gRPC插件的代码，可以发现generator.RegisterPlugin函数可以用来注册插件。插件是一个generator.Plugin接口：
// A Plugin provides functionality to add to the output during
// Go code generation, such as to produce RPC stubs.
type Plugin interface {
// Name identifies the plugin.
Name() string
// Init is called once after data structures are built but before
// code generation begins.
Init(g *Generator)
// Generate produces the code generated by the plugin for this file,
// except for the imports, by calling the generator's methods P, In,
// and Out.
Generate(file *FileDescriptor)
// GenerateImports produces the import declarations for this file.
// It is called after Generate.
GenerateImports(file *FileDescriptor)
}其中Name方法返回插件的名字，这是Go语言的Protobuf实现的插件体系，和protoc插件的名字并无关系。然后Init函数是通过g参数对插件进行初始化，g参数中包含Proto文件的所有信息。最后的Generate和GenerateImports方法用于生成主体代码和对应的导入包代码。
因此我们可以设计一个netrpcPlugin插件，用于为标准库的RPC框架生成代码：
import (
"github.com/golang/protobuf/protoc-gen-go/generator"
)
type netrpcPlugin struct{ *generator.Generator }
func (p *netrpcPlugin) Name() string { return "netrpc" }
func (p *netrpcPlugin) Init(g *generator.Generator) { p.Generator = g }
func (p *netrpcPlugin) GenerateImports(file *generator.FileDescriptor) {
if len(file.Service) > 0 {
p.genImportCode(file)
}
}
func (p *netrpcPlugin) Generate(file *generator.FileDescriptor) {
for _, svc := range file.Service {
p.genServiceCode(svc)
}
}首先Name方法返回插件的名字。netrpcPlugin插件内置了一个匿名的*generator.Generator成员，然后在Init初始化的时候用参数g进行初始化，因此插件是从g参数对象继承了全部的公有方法。其中GenerateImports方法调用自定义的genImportCode函数生成导入代码。Generate方法调用自定义的genServiceCode方法生成每个服务的代码。
目前，自定义的genImportCode和genServiceCode方法只是输出一行简单的注释：
func (p *netrpcPlugin) genImportCode(file *generator.FileDescriptor) {
p.P("// TODO: import code")
}
func (p *netrpcPlugin) genServiceCode(svc *descriptor.ServiceDescriptorProto) {
p.P("// TODO: service code, Name = " + svc.GetName())
}要使用该插件需要先通过generator.RegisterPlugin函数注册插件，可以在init函数中完成：
func init() {
generator.RegisterPlugin(new(netrpcPlugin))
}因为Go语言的包只能静态导入，我们无法向已经安装的protoc-gen-go添加我们新编写的插件。我们将重新克隆protoc-gen-go对应的main函数：
package main
import (
"io/ioutil"
"os"
"github.com/golang/protobuf/proto"
"github.com/golang/protobuf/protoc-gen-go/generator"
)
func main() {
g := generator.New()
data, err := ioutil.ReadAll(os.Stdin)
if err != nil {
g.Error(err, "reading input")
}
if err := proto.Unmarshal(data, g.Request); err != nil {
g.Error(err, "parsing input proto")
}
if len(g.Request.FileToGenerate) == 0 {
g.Fail("no files to generate")
}
g.CommandLineParameters(g.Request.GetParameter())
// Create a wrapped version of the Descriptors and EnumDescriptors that
// point to the file that defines them.
g.WrapTypes()
g.SetPackageNames()
g.BuildTypeNameMap()
g.GenerateAllFiles()
// Send back the results.
data, err = proto.Marshal(g.Response)
if err != nil {
g.Error(err, "failed to marshal output proto")
}
_, err = os.Stdout.Write(data)
if err != nil {
g.Error(err, "failed to write output proto")
}
}为了避免对protoc-gen-go插件造成干扰，我们将我们的可执行程序命名为protoc-gen-go-netrpc，表示包含了netrpc插件。然后用以下命令重新编译hello.proto文件：
$ protoc --go-netrpc_out=plugins=netrpc:. hello.proto其中--go-netrpc_out参数告知protoc编译器加载名为protoc-gen-go-netrpc的插件，插件中的plugins=netrpc指示启用内部唯一的名为netrpc的netrpcPlugin插件。在新生成的hello.pb.go文件中将包含增加的注释代码。
至此，手工定制的Protobuf代码生成插件终于可以工作了。
学员评价