protoc go插件编写之四 (实现生成自己的proto文件)

前面几篇文章讲述了protoc 的原理，实现，及参数，这里终于要开始动手了。

准备工作

• 1. 安装go 1.14
• 2. 在一个新的目录下, 建立一个go项目

go mod init github.com/kekek/protoc-gen-foo
go get google.golang.org/protobuf@v1.21.0

⚠️： 这里文件夹的名字及其重要，必须以protoc-gen- 开头， 比如proto-gen-foo

plan

在此示例中，我们想扩展Protobuf编译器以添加方法Foo（），该方法为.proto文件中定义的每条消息返回bar。

建立如下test.proto 文件

syntax = "proto3";

package test;

option go_package = ".;test";
message Message {
    string data = 1;
}

调用编译器插件

首先，我们将编写一个无用的插件，该插件将写入stdout。将以下内容添加到main.go中：

package main

import (
   "log"
)

func main()  {
   log.Println("hello protoc")
   return
}

运行如下代码，调用插件：

# 建立文件夹out，输出生成内容
mkdir out

# 安装插件
go install .

# 运行 protoc，传入 --foo_out 参数
protoc \
    --proto_path . \
    -I=. \
    test.proto \
    --foo_out=./out \
    --go_out=./out

⚠️ protoc 通过 --foo_out 搜索插件 可执行文件 protoc-gen-foo， 也可使用参数 protoc --plugin=protoc-gen-foo=/path/to/protoc-gen-foo 指定插件位置

如果上述命令成功执行，则应观察到hello protoc已打印到控制台，以及生成的protobuf文件out / test.pb.go。

生成代码

用以下代码替换我们的准备系统插件，请注意注释中的解释：

package main

import (
   "bytes"
   "fmt"
   "google.golang.org/protobuf/compiler/protogen"
   "google.golang.org/protobuf/proto"
   "google.golang.org/protobuf/types/pluginpb"
   "io/ioutil"
   "os"
)

func main()  {

   // Protoc 将protobuf文件编译为 pluginpb.CodeGeneratorRequest结构，并输出到stdin中
   input, _ := ioutil.ReadAll(os.Stdin)
   var req pluginpb.CodeGeneratorRequest
   proto.Unmarshal(input, &req)

   // 使用默认选项初始化我们的插件
   opts := protogen.Options{}
   plugin, err := opts.New(&req)
   if err != nil {
      panic(err)
   }

   // protoc 将一组文件结构传递给程序处理
   for _, file := range plugin.Files {

      // 是时候生成代码了……！

      // 1. 初始化缓冲区以保存生成的代码
      var buf bytes.Buffer

      // 2. 生成包名称
      pkg := fmt.Sprintf("package %s", file.GoPackageName)
      buf.Write([]byte(pkg))

      // 3. 为每个message生成 Foo() 方法
      for _, msg := range file.Proto.MessageType {
         buf.Write([]byte(fmt.Sprintf(`
            func (x %s) Foo() string {
               return "bar"
            }`, *msg.Name)))
      }

      // 4. 指定输出文件名，在这种情况下为test.foo.go
      filename := file.GeneratedFilenamePrefix + ".foo.go"
      file := plugin.NewGeneratedFile(filename, ".")

      // 5. 将设概念车呢个的代码，从缓冲区写入到文件
      file.Write(buf.Bytes())
   }

   // 从我们的插件生成响应,并将其编组为protobuf
   stdout := plugin.Response()
   out, err := proto.Marshal(stdout)
   if err != nil {
      panic(err)
   }

   // 相应输出到stdout, 它将被 protoc 接收
   fmt.Fprintf(os.Stdout, string(out))
}

重新运行上一部分中的命令，并观察新生成的文件out/test.foo.go

package test

func (x Message) Foo() string {
   return "bar"
}

一个简单的protocbuf 完成了

• 本文翻译自 https://medium.com/@tim.r.coulson/writing-a-protoc-plugin-with-google-golang-org-protobuf-cd5aa75f5777

有疑问加站长微信联系（非本文作者）