golang源码分析：自定义proto插件
golang源码分析：自定义proto插件
在使用protoc的时候，可以通过指定不同的插件来生成不同的代码，它的参数统一是xx_out结尾的，制定了插件参数，就会到path下搜索protoc-gen-xx的插件。比如 protoc 通过 --foo_out 搜索插件 可执行文件 protoc-gen-foo， 也可使用参数 protoc --plugin=protoc-gen-foo=/path/to/protoc-gen-foo 指定插件位置。
protoc插件是一个独立的二进制程序，protoc进程通过fork生成子进程，并exec加载插件程序运行。父子进程间通过管道通信，并将管道的输入和输出重定向到标准输入和标准输出。
protoc进程将proto文件的信息封装为CodeGeneratorRequest传递给插件子进程，插件子进程将根据CodeGeneratorRequest中的信息，将要生成的代码数据封装为CodeGeneratorResponse对象传递给protoc进程。
插件进程从标准输入读取出CodeGeneratorRequest数据，将CodeGeneratorResponse数据写到标准输出。CodeGeneratorRequest和CodeGeneratorRequest两者也是使用proto定义的。于是一个protoc插件的开发可以简单分为三步：
- 从标准输入读取解析出CodeGeneratorRequest数据
- 利用读取的数据来生成对应的代码
- 将生成的结果封装为CodeGeneratorResponse写入标准输出
编写protoc插件的模板代码如下：
package main
import (
"flag"
"fmt"
"google.golang.org/protobuf/compiler/protogen"
)
func main() {
// 用于接收命令行参数
var (
flags flag.FlagSet
plugins = flags.String("plugins", "", "list of plugins to enable (supported values: grpc)")
importPrefix = flags.String("import_prefix", "", "prefix to prepend to import paths")
)
importRewriteFunc := func(importPath protogen.GoImportPath) protogen.GoImportPath {
switch importPath {
case "context", "fmt", "math":
return importPath
}
if *importPrefix != "" {
return protogen.GoImportPath(*importPrefix) + importPath
}
return importPath
}
protogen.Options{
ParamFunc: flags.Set,
ImportRewriteFunc: importRewriteFunc,
}.Run(func(gen *protogen.Plugin) error {
// ...
for _, f := range gen.Files {
// 根据proto文件信息来生成新文件
fmt.Println(plugins, f)
}
return nil
})
}其中
protogen.Options{
ParamFunc: flags.Set,
ImportRewriteFunc: importRewriteFunc,
}Options有两个字段：
- ParamFunc：命令行中的插件参数会以--go_out=<param1>=<value1>,<param2>=<value2>:<output_directory>的形式输入并最总被解析为CodeGeneratorRequest字段，Run方法运行过程中会读取出键值对并调用ParamFunc函数，这样就可以将命令行参数绑定到flags对应的变量中了。
- ImportRewriteFunc：生成的新文件中的每个包导入的路径可以使用此函数进行重写
然后调用Run方法来进行相关代码的生成。下面我们实现一个简单的插件，实现解析消息体的字段名，并写文件。
package main
import (
"strconv"
"strings"
"google.golang.org/protobuf/compiler/protogen"
)
func main() {
protogen.Options{}.Run(func(p *protogen.Plugin) error {
// 遍历proto文件
for _, f := range p.Files {
fname := f.GeneratedFilenamePrefix + ".txt"
// 后续使用t来写入新文件
t := p.NewGeneratedFile(fname, f.GoImportPath)
for _, msg := range f.Messages {
builder := strings.Builder{}
for _, field := range msg.Fields {
builder.WriteString(field.Desc.TextName() + ": " + strconv.Itoa(field.Desc.Index()) + "\n")
}
t.Write([]byte(builder.String()))
}
}
return nil
})
}可以定义一个文件来测试下
syntax = "proto3";
package api;
option go_package = "api/v1;v1";
message HelloRequest {
string msg = 1;
} export PATH=$GOPATH/bin:$PATH
go build -o protoc-gen-test main.go
cp protoc-gen-test $GOPATH/bin
chmod +X $GOPATH/bin/proto-gen-test然后就可以使用我们自定义的插件了
protoc --test_out=. test.proto可以看到，生成了文件test.txt,内容是
msg: 0至此一个简单的protoc插件开发完毕。
本文分享自 golang算法架构leetcode技术php 微信公众号，前往查看
如有侵权，请联系 cloudcommunity@tencent.com 删除。
本文参与 腾讯云自媒体同步曝光计划 ，欢迎热爱写作的你一起参与！
腾讯云开发者
