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 插件的開發可以簡單分爲三步：

  1. 從標準輸入讀取解析出 CodeGeneratorRequest 數據

  2. 利用讀取的數據來生成對應的代碼

  3. 將生成的結果封裝爲 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 有兩個字段：

  1. ParamFunc：命令行中的插件參數會以 --go_out==,=:<output_directory> 的形式輸入並最總被解析爲 CodeGeneratorRequest 字段，Run 方法運行過程中會讀取出鍵值對並調用 ParamFunc 函數，這樣就可以將命令行參數綁定到 flags 對應的變量中了。

  2. 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 插件開發完畢。

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/kS2EJt1V6kStm8CgprslCg

猜你喜歡