0%

protoc插件编写

利用protoc,可以解析protobuf文件生成。一般来说都是通过自定义插件,调用库方法来实现。官方的protobuf解析实际上也是利用插件完成的,例如protoc-gen-go

早先使用c++版本书写的时候需要从stdin里获取pb文件的数据,之后再根据descriptor生成代码片段(或者是别的文本片段,txt,json或者别的你需要的文本,根据需求而定),最后把文本输出到stdout里。我们可以理解为实际上protoc就是把proto文件通过标准输入喂给了插件,之后截取标准输出作为生成的文件。这里实际上步骤是很繁琐的,所以之前很多同学宁可用点第三方的parser

不过官方后续给go提供了一个比较好使的库,google.golang.org/protobuf/compiler/protogen,因此后续写插件我都推荐用go加上这个库来写。虽然用的是go,但是写出来的插件可以生成各种代码,只要是文本格式就可以(事实上不是文本也毫无问题)。

最简单的插件

先从一个最简单的插件开始:根据protobuf文件输出它的package namemessage列表,以及rpc列表到一个文本文件:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
syntax = "proto3";
package test;

// Foo ...
message Foo {
int32 a = 1;
string b = 2;
}

message Bar {
int64 a = 1;
}

service TestService {
rpc Test (Foo) returns (Bar) {}
}

新建一个go项目,因为protoc的插件名字约定为protoc-gen-xxx,所以最好就叫做protoc-gen-txt

插件的整体代码非常简单:

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
26
27
28
29
30
31
32
package main

import (
"fmt"
"google.golang.org/protobuf/compiler/protogen"
)

func main() {
protogen.Options{}.Run(
func(p *protogen.Plugin) error {
for _, f := range p.Files {
if f.Generate {
filename := f.GeneratedFilenamePrefix + ".txt"
g := p.NewGeneratedFile(filename, f.GoImportPath)
g.Write([]byte(fmt.Sprintf("package name: %s\n", *f.Proto.Package)))
g.Write([]byte("Messages: \n"))
for _, msg := range f.Messages {
g.Write([]byte(" -"))
g.Write([]byte(msg.Desc.Name()))
g.Write([]byte("\n"))
}
for _, service := range f.Services {
g.Write([]byte(fmt.Sprintf("service : %s\n", service.Desc.Name())))
for _, rpc := range service.Methods {
g.Write(([]byte(fmt.Sprintf(" - %s\n", rpc.Desc.Name()))))
}
}
}
}
return nil
})
}

先编译之后在项目目录下运行 protoc --plugin ./protoc-gen-txt --txt_opt=Mtest.proto=./ --txt_out=./ ./test.proto,可以看到一个新生成了一个test.txt文件:

1
2
3
4
5
6
package name: test
Messages:
-Foo
-Bar
service : TestService
- Test

插件运行指令分析

首先我们分析一下执行的指令,protoc肯定没话说,整个转换的主函数,装了protoc就有的工具。

--plugin ./protoc-gen-txt是告知protoc,这次要加载一个叫做protoc-gen-txt的插件,它的路径就在当前目录下。当然如果这个插件本身就在系统路径里,也可以不加。

--txt_opt=Mtest.proto=./前半段代表一个选项参数,txt_opt代表给protoc-gen-txt这个插件传递的参数(参数我们后续再聊)。当然如果你是要给protoc-gen-go这个插件传递参数,那么就可以设置成--go_opt=xxx。后半段代表test.protogo import path.,通过M+文件名来声明。这是为了让插件能顺利执行,本身并没有特别意义。

--txt_out=./用于指定输出目录,即txt文件的输出目录。

插件代码分析

接下来简要分析一下插件代码:

这里首先生成了一个protogen.Options,样例中我们传的是空,实际上可以传递命令行参数进来,后续我们再考虑。这个optionsRun方法接受一个func(*protogen.Plugin) error形式的函数或者方法作为参数。

其中protogen.Plugin是一个protoc插件的调用操作,可以看到其中包含了这次解析的文件的各种信息,我们需要的内容都在这个结构体里包含了。

之后的NewGeneratedFile生成了一个输出文件句柄,并且获得了它的引用。之后我们就可以调用Write方法把文件内容输出了。至于文件的flushclose方法没有提供,也无需我们手动调用。如果对于一个proto文件我们需要输出两个文件,可以多调用NewGeneratedFile一次获得一个新的文件句柄,参考如下修改:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
filename := f.GeneratedFilenamePrefix + ".txt"
g := p.NewGeneratedFile(filename, f.GoImportPath)
g2 := p.NewGeneratedFile(f.GeneratedFilenamePrefix+"2.txt", f.GoImportPath) // 新增一个文件
g.Write([]byte(fmt.Sprintf("package name: %s\n", *f.Proto.Package)))
g.Write([]byte("Messages: \n"))
for _, msg := range f.Messages {
g.Write([]byte(" -"))
g.Write([]byte(msg.Desc.Name()))
g.Write([]byte("\n"))
}
// 以下关于service的信息都输出到了另一个test2.txt中
for _, service := range f.Services {
g2.Write([]byte(fmt.Sprintf("service : %s\n", service.Desc.Name())))
for _, rpc := range service.Methods {
g2.Write(([]byte(fmt.Sprintf(" - %s\n", rpc.Desc.Name()))))
}
}

我们会得到两个文件,更多的文件需求以此类推,不再赘述。

protobuf 文件结构

如样例中所示,我们需要的所有有关于protobuf文件的信息都存储在protogen.File中。正常情况下我们常用的结构体有这些:

  • protogen.Message 用于描述Message的定义
  • protogen.Serviceprotogen.Method用于描述定义的service和它包含的rpc的信息
  • protoreflect.Descriptor定义了一堆的描述,用于描述一个消息/方法/服务的名字。例如protoreflect.MessageDescriptor用于描述messageprotoreflect.FieldDescriptor用于描述message的每个field

因为库代码是开源的,因此上述的结构体描述都可以在源文件中看到,按需取用即可。

使用模板生成文本

在上一节中我们尝试用tab来让生成的文件带有点格式化,显然这么做十分低效,而且非常不直观,稍不留意就可能搞坏格式(此时假设你需要生成的是一份python代码……)。这里介绍一个比较好用的包:text/template,方便我们使用模板来生成一些通用的文本文件,例如 HTML 文件和配置文件。

首先先建立一个模版,out.tmpl

1
2
3
4
5
6
7
8
9
10
11
12
13
14
{{if .Generate}}
package name: {{.Proto.Package}}
Messages:
{{- range .Messages }}
- {{.Desc.Name}}
{{- end }}

{{- range .Services }}
service : {{ .Desc.Name }}
{{- range .Methods }}
- {{ .Desc.Name }}
{{- end }}
{{- end }}
{{end}}

之后将我们的代码修改成这样:

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
26
27
28
29
package main

import (
"embed"
"text/template"

"google.golang.org/protobuf/compiler/protogen"
)

//go:embed out.tmpl
var outTemplate embed.FS

func main() {
protogen.Options{}.Run(
func(p *protogen.Plugin) error {
for _, f := range p.Files {
filename := f.GeneratedFilenamePrefix + ".txt"
g := p.NewGeneratedFile(filename, f.GoImportPath)

t := template.New("out")

tfile := template.Must(t.ParseFS(outTemplate, "*"))
if err := tfile.Lookup("out.tmpl").Execute(g, f); err != nil {
panic(err)
}
}
return nil
})
}

这里引入了几个库,首先是embed

它可以将文件或一个或多个文件的内容嵌入到 Go 二进制文件中,允许在程序运行时访问这些文件的内容,而无需读取外部文件。因此,这个包使得分发静态资源变得更加方便。

因此我们利用它将out.tmpl打包到二进制文件中,否则我们的插件还会依赖模版文件,就相当不妙。

1
2
3
// 注意这边这个注释是go的编译指令,不可以省略
//go:embed out.tmpl
var outTemplate embed.FS

其次是template包:

1
2
3
4
5
tfile := template.Must(t.ParseFS(outTemplate, "*"))
// 这里的g是需要生成的文件 f是*protogen.File
if err := tfile.Lookup("out.tmpl").Execute(g, f); err != nil {
panic(err)
}

相当于将protobuf的文件信息丢给这个模板,最终期望将数据写入文件中。

template语法

首先是变量,在模版中,变量用{{ }}括起来:

1
package name: {{.Proto.Package}}

其中.Proto.Package是一个操作符,代表当前模板中的上下文的.Proto.Package。此时模板的上下文是f (*protogen.File)(在代码中传入的),因此这里可以认为是

1
package name: {{f.Proto.Package}}

和我们之前的代码很像:

1
g.Write([]byte(fmt.Sprintf("package name: %s\n", *f.Proto.Package)))

其次是循环

1
2
3
{{- range .Messages }}
- {{.Desc.Name}}
{{- end }}

很显然这里的.Message就代表了f.Message{{- range .Messages }}则代表了遍历f.Message

但是需要注意的是,循环中的上下文已经不是f了,而是遍历.Messages时当前迭代的元素。这里的效果其实也等效于之前的代码:

1
2
3
4
5
for _, msg := range f.Messages {
g.Write([]byte(" -"))
g.Write([]byte(msg.Desc.Name()))
g.Write([]byte("\n"))
}

最后是判断

模版的开头就使用以下语句根据f.Gengerate判断是否要生成对应的txt文件:

1
2
3
{{if .Generate}}
...
{{end}}

有点类似原来代码里的这段:

1
2
3
if f.Generate {
....
}

但是有稍许不同,因为原来的代码不会走到文件生成这里,但是在模板里判断无论如何都会生成一个txt文件。

模板的进阶用法

在前一个小节里我们看到,模板可以直接引用到上下文的数据信息。但是有时候我们需要输出的信息可能会需要插件额外加工,例如我们可能需要获取全小写的消息名,那么就不能直接引用原始的数据,而是希望能够调用我们自定义的方法。为此我们需要对模版做一些修改:

1
2
3
4
5
package name: {{.Proto.Package}}
Messages:
{{- range .Messages }}
- {{.Desc.Name}}, lower case: {{LowerName .Desc.Name}}
{{- end }}

请注意,这里的LowerName并不以.开头了,这在模板的语法中代表需要调用一个外部的函数。而调用的参数则是.Desc.Name。此时我们对插件的代码也做出一些修改:

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
26
27
28
29
30
31
32
33
34
import (
"embed"
"strings"
"text/template"

"google.golang.org/protobuf/compiler/protogen"
"google.golang.org/protobuf/reflect/protoreflect"
)

//go:embed out.tmpl
var outTemplate embed.FS

func main() {
protogen.Options{}.Run(
func(p *protogen.Plugin) error {
for _, f := range p.Files {
filename := f.GeneratedFilenamePrefix + ".txt"
g := p.NewGeneratedFile(filename, f.GoImportPath)

t := template.New("out")
// 请注意这里 新增了一个映射
t.Funcs(map[string]interface{}{
"LowerName": func(name protoreflect.Name) string {
return strings.ToLower(string(name))
},
})
tfile := template.Must(t.ParseFS(outTemplate, "*"))
if err := tfile.Lookup("out.tmpl").Execute(g, f); err != nil {
panic(err)
}
}
return nil
})
}

和的不同就是新增了这一段:

1
2
3
4
5
t.Funcs(map[string]interface{}{
"LowerName": func(name protoreflect.Name) string {
return strings.ToLower(string(name))
},
})

很显然,这里是通过增加了一个标识符到匿名函数的映射。当模板解析到LowerName时,就知道调用对应的函数,并且将返回值用于替换变量。

插件的命令行参数

在第一章中我们跑第一个例子就见到了插件的命令行参数了:

1
protoc --plugin ./protoc-gen-txt --txt_opt=Mtest.proto=./ --txt_out=./ ./test.proto

这里的--txt_opt=Mtest.proto=./就是命令行参数,txt_opt代表给protoc-gen-txt的参数。我们可以对原有的插件代码做一些修改,让给插件传递一些参数:

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
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
package main

import (
"embed"
"flag"
"strings"
"text/template"

"google.golang.org/protobuf/compiler/protogen"
"google.golang.org/protobuf/reflect/protoreflect"
)

//go:embed out.tmpl
var outTemplate embed.FS

func main() {
var flags flag.FlagSet

// 这里增加了一个flag
needGenerate := flags.Bool("generate", false, "if we need generate")
protogen.Options{
ParamFunc: flags.Set,
}.Run(
func(p *protogen.Plugin) error {
for _, f := range p.Files {
if !*needGenerate {
continue
}
filename := f.GeneratedFilenamePrefix + ".txt"
g := p.NewGeneratedFile(filename, f.GoImportPath)

t := template.New("out")
t.Funcs(map[string]interface{}{
"LowerName": func(name protoreflect.Name) string {
return strings.ToLower(string(name))
},
})
tfile := template.Must(t.ParseFS(outTemplate, "*"))
if err := tfile.Lookup("out.tmpl").Execute(g, f); err != nil {
panic(err)
}
}
return nil
})
}

使用以下命令行参数进行转表

1
protoc --plugin ./protoc-gen-txt --txt_out=./ --txt_opt=Mtest.proto=./,generate=true ./test.proto 

或者也可以使用如下方式将命令行参数加入进来:

1
protoc --plugin ./protoc-gen-txt --txt_out=Mtest.proto=./,generate=true:. ./test.proto

读取扩展信息

通常我们在业务上会需要定义一些扩展信息方便使用,先增添一个option.proto文件:

1
2
3
4
5
6
7
8
9
syntax = "proto2";
import "google/protobuf/descriptor.proto";

package option_test;
option go_package = "./optiontest";

extend google.protobuf.MessageOptions {
optional bool needgenerate = 10001; // 是否需要生成
}

之后修改原有的test.proto

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
syntax = "proto3";
package test;

import "option.proto";

// Foo ...
message Foo {
option (option_test.needgenerate) = true;
int32 a = 1;
string b = 2;
}

message Bar { int64 a = 1; }

service TestService {
rpc Test(Foo) returns (Bar) {}
}

这样就引入了一个扩展,用于标记是否要生成这个message的信息,为了利用这个信息,我们可以对main.goout.tmpl稍作修改:

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
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
package main

import (
"embed"
"protoc-gen-txt/optiontest"
"strings"
"text/template"

"google.golang.org/protobuf/compiler/protogen"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/reflect/protoreflect"
)

//go:embed out.tmpl
var outTemplate embed.FS

func main() {
protogen.Options{}.Run(
func(p *protogen.Plugin) error {
for _, f := range p.Files {
if !f.Generate {
continue
}
filename := f.GeneratedFilenamePrefix + ".txt"
g := p.NewGeneratedFile(filename, f.GoImportPath)

t := template.New("out")
// 请注意这里 新增了一个映射
t.Funcs(map[string]interface{}{
"LowerName": func(name protoreflect.Name) string {
return strings.ToLower(string(name))
},
"Generate": func(desc protoreflect.MessageDescriptor) bool {
generate, ok := proto.GetExtension(desc.Options(), optiontest.E_Needgenerate).(bool) // 在这里读取扩展
if !ok {
return false
}
return generate
},
})
tfile := template.Must(t.ParseFS(outTemplate, "*"))
if err := tfile.Lookup("out.tmpl").Execute(g, f); err != nil {
panic(err)
}
}
return nil
})
}

1
2
3
4
5
6
7
package name: {{.Proto.Package}}
Messages:
{{- range .Messages }}
{{if Generate .Desc}}
- {{.Desc.Name}}, lower case: {{LowerName .Desc.Name}}
{{end}}
{{- end }}

为了编译这个工具,这次要多加一步:

1
2
protoc --go_out=./ ./option.proto
protoc --plugin ./protoc-gen-txt --txt_opt=Mtest.proto=./ --txt_out=./ ./test.proto

第一步用于将option.proto编译成go代码,在main.go中引用。

分析代码,在模板里多了这么一段,用于判断是否需要生成对应msg的信息:

1
2
3
{{if Generate .Desc}}
- {{.Desc.Name}}, lower case: {{LowerName .Desc.Name}}
{{end}}

根据模板里的说明,我们知道这里会试图寻找一个名为Generate的函数,它的实现在main.go里:

1
2
3
4
5
6
7
"Generate": func(desc protoreflect.MessageDescriptor) bool {
generate, ok := proto.GetExtension(desc.Options(), optiontest.E_Needgenerate).(bool) // 在这里读取扩展
if !ok {
return false
}
return generate
},

这里的核心逻辑就在于proto.GetExtension,这是用于获取Msg的扩展用的。其中optiontest.E_Needgenerate就对应了option.proto里的needgenerate。生成的代码在option.pb.go里,代码不长,可以看看扩展在go中是怎么定义的。

获取了指定的扩展信息之后,就可以读出在msg中扩展的赋值,这里我们用一个bool值判断是否需要生成代码。之后如果在业务上大家也有类似的需要,也可以使用别的值,例如int32string