ml の 小破绽

如何编写一个 protoc 的插件

我们在编写 protocol 文件去生成.pb.go 文件的时候你都知道它是如何做到的吗？
整个过程很简单，其实我们之前生成的 go 文件是用的 protoc-gen-go 插件，可以去 github 上看看这个插件的代码 https://github.com/protocolbuffers/protobuf-go/blob/master/cmd/protoc-gen-go/main.go 里面主要用到的库是 google.golang.org/protobuf/compiler/protogen

其实插件的工作原理就是从标准输出读取和输入,具体可以看 https://protobuf.dev/reference/other/ 然后根据你的命令行中的 –{NAME}_out 去找 protoc-gen-{NAME} 这个插件然后生成文件具体请看官方文档
今天我们就一起来实现一个插件吧！

首先我们先来生成一个最基本的.pb.go 文件：

1. 在生成之前我们都需要去下载一个二进制文件

   go install google.golang.org/protobuf/cmd/protoc-gen-go@latest

   这个会存在你的$GOPATH/bin 目录下

2. 编写一个简单的 proto

   syntax = "proto3";
   package simple;
   option go_package = "./simple";
   
   message hello {
     string world = 1;
   }

3. 生成 .pb.go 文件
   为了后续方便，就直接写一个 makefile

   gen:
       protoc --go_out=. --go_opt=paths=source_relative ./simple/simple.proto

   执行 make gen 就能输出 .pb.go 文件了

   

插件编写

1. 代码

   package main
   
   import (
     "flag"
     "fmt"
     "google.golang.org/protobuf/compiler/protogen"
     "os"
   )
   
   func main() {
     protogen.Options{
       ParamFunc: flag.CommandLine.Set,
     }.Run(func(gen *protogen.Plugin) error {
       for _, f := range gen.Files {
         if f.Generate {
           GenerateFile(gen, f)
         }
       }
   
       return nil
     })
   }
   
   func GenerateFile(gen *protogen.Plugin, f *protogen.File) {
     filename := f.GeneratedFilenamePrefix + "_plugin.pb.go"
     g := gen.NewGeneratedFile(filename, f.GoImportPath) // 初始化一个需要生成的文件
     g.P("// Code generated by protoc-gen-simple. DO NOT EDIT.")
     g.P("// source: ", f.Desc.Path())
     g.P()
     g.P("package ", f.GoPackageName) //包名
     for _, message := range f.Messages {
       _, _ = fmt.Fprint(os.Stderr, message)   // 主要用来调试
       g.P("var _ = &", message.GoIdent, "{}") // 输出
     }
   }

   • gen.NewGeneratedFile 实例化一个生成文件的结构体

   • g.P(…) 写入文件操作

   • f.Messages 就是里面的 message 定义 还有 service 等等

   • _, _ = fmt.Fprint(os.Stderr, message)

     至于这一段是为了调试，由于需要编译无法进行单步调试我就将需要获取的调试信息打印在了os.Stderr ,至于为啥不能是标准输出呢，是因为插件就是用标准输出来通信的，前面也有讲到过

2. 编译

   go install

   然后你就能看到你的$GOPATH/bin 目录下会有一个 protoc-gen-go-simple 二进制文件

3. 使用插件
   在之前的命令加上一句 –go-simple_out=. –go-simple_opt=paths=source_relative

   gen:
     protoc \
         --go_out=. --go_opt=paths=source_relative \
         --go-simple_out=. --go-simple_opt=paths=source_relative \
         ./simple/simple.proto

   make gen 你就能看到有两个.pb.go 文件了
   

以上的过程还是很简单的，就是调试起来很费劲也可能还有很简单的调试办法我不知道，还有一点就是文档不太友好，还得去看源码。基本上 proto 文件中所定义的里面的 *protogen.File 都能够拿到。所以可以干很多的骚操作，赶紧学起来吧～