Go で作る protoc プラグイン入門

はじめに

本記事では、Go 言語で protoc プラグインを作成して実際に動かす方法を紹介します。

ここで作るプラグインの動作を紹介します。まず .proto ファイルに次のような message が定義されているとします。

message Person {
  string name = 1;
  int32 age = 2;
  Company company = 3;
}

message Company {
  string name = 1;
}

プラグインでコードを生成することで、次のように Person 構造体と Company 構造体を利用できるようになり、また YAML 形式で出力できるようになります。

func main() {
  p := proto.Person{
    Name: "Taro",
    Age:  30,
    Company: proto.Company{
      Name: "FooBarCorp",
    },
  }

  p.PrintYAML(0)
}

実行結果は次の通りです。

$ go run main.go
Name: Taro
Age: 30
Company:
  Name: FooBarCorp

本記事ではプラグインを実際に動かすための完全なコードを記載しますが、Protocol Buffers や protoc プラグインの概要についてはすでに良い記事があるため詳細を省きます。これらの概要を知りたい方には yugui さんの記事をおすすめします。

protocプラグインの書き方 - Qiita
以前の記事では、Protocol Buffers (protobuf)の魅力の1つは周辺ツールを拡張しやすいことだと述べた。そこで本稿では具体的に拡張のためのprotocプラグインの書き方を紹介したい。 ちなみに、protobufの周辺ツールと言うと2種類ある。 1つはp…

準備

事前に go コンパイラと protoc がインストールされている必要があります。

次のようにディレクトリとファイルを用意します。

go mod init example.com/protoc-gen-subgo
mkdir -p bin cmd/protoc-gen-subgo protobuf/go protobuf/proto
touch cmd/protoc-gen-subgo/main.go
touch protobuf/proto/person.proto
touch main.go

これで次のディレクトリ構造になります。

$ tree -F
./
├── bin/
├── cmd/
│   └── protoc-gen-subgo/
│       └── main.go
├── go.mod
├── main.go
└── protobuf/
    ├── go/
    └── proto/
        └── person.proto

Protocol Buffers Message の作成

プラグインで処理するための message を protobuf/proto/person.proto に書きます。

syntax = "proto3";

package example;
option go_package = "example.com/protoc-gen-subgo/protobuf/proto";

message Person {
  string name = 1;
  int32 age = 2;
  Company company = 3;
}

message Company {
  string name = 1;
}

protobuf/proto/person.proto の内容

protoc プラグインによる構造体の生成

Go 言語で protoc プラグインを実装していきます。プラグインの実装にあたっては protogen パッケージを利用します。

まずは各 message に対応する構造体を生成できるようにします。cmd/protoc-gen-subgo/main.go に次の内容を書きます。

package main

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

func main() {
  opts := protogen.Options{}
  opts.Run(func(plugin *protogen.Plugin) error {
    for _, f := range plugin.Files {
      fn := f.GeneratedFilenamePrefix + ".sub.go"
      g := plugin.NewGeneratedFile(fn, f.GoImportPath)
      g.P("package ", f.GoPackageName)

      for _, m := range f.Messages {
        err := processMessage(g, m)
        if err != nil {
          return err
        }
      }
    }

    return nil
  })
}

func processMessage(g *protogen.GeneratedFile, m *protogen.Message) error {
  g.P("type ", m.GoIdent, " struct {")
  g.P("}")

  return nil
}
  • Options.Run に引数として与えた関数が protoc プラグインとして実行されます。
  • Plugin.NewGeneratedFile によって *GeneratedFile 型の値が得られ、GeneratedFile.P メソッドを呼び出すことで生成コードを書きこむことができます。例えば g.P("// TODO") を実行すると、生成されるコードには // TODO という内容が追記されることになります。

早速プラグインでコードを生成してみます。次のように bin/protoc-gen-subgo をビルドします。

go mod tidy
go build -o ./bin/protoc-gen-subgo ./cmd/protoc-gen-subgo

protoc は実行時のオプション名に従って対応する protoc プラグインを呼び出すようになっており、protoc --subgo_out=./some/path ...のように実行されたときは protoc-gen-subgo という名前のファイルを実行するようになっています。そこで、./bin を PATH 環境変数に追加することで protoc-gen-subgo が実行されるようにします。

export PATH=$PATH:$PWD/bin

実際に protoc からプラグインを利用します。

protoc -I protobuf/proto \
  --subgo_out=paths=source_relative:./protobuf/go \
  ./protobuf/proto/person.proto

これによって次のような protobuf/go/person.sub.go が生成されます。

package person

type Person struct {
}
type Company struct {
}

protobuf/go/person.sub.go

続いて、message の各フィールドに対応するフィールドが struct に追加されるようにします。cmd/protoc-gen-subgo/main.go を次のように変更します。

import (
	"fmt" // 追加

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

func processMessage(g *protogen.GeneratedFile, m *protogen.Message) error {
  g.P("type ", m.GoIdent, " struct {")
  
  // 次を追加
  for _, f := range m.Fields {
    ty, err := goTypeName(f)
    if err != nil {
      return err
    }
    g.P(f.GoName, " ", ty)
  }
  
  g.P("}")

  return nil
}

// 追加
func goTypeName(f *protogen.Field) (string, error) {
  switch k := f.Desc.Kind(); k {
  case protoreflect.Int32Kind:
    return "int32", nil
  case protoreflect.StringKind:
    return "string", nil
  case protoreflect.MessageKind:
    return f.Message.GoIdent.GoName, nil
  default:
    return "", fmt.Errorf("未対応の Field: %s %v", f.GoName, k)
  }
}

再度 protoc-gen-subgo をビルドした上で protoc を実行すると、protobuf/go/person.sub.go は次のようになります。

package person

type Person struct {
  Name    string
  Age     int32
  Company Company
}
type Company struct {
  Name string
}

この構造体は次のように ./main.go から利用することができます。

package main

import (
  "fmt"

  proto "example.com/protoc-gen-subgo/protobuf/go"
)

func main() {
  p := proto.Person{
    Name: "Taro",
    Age:  30,
    Company: proto.Company{
      Name: "FooBarCorp",
    },
  }
  fmt.Printf("%v\n", p)
}

./main.go の内容

ここまでの実装により、go run main.go を実行すると {Taro 30 {FooBarCorp}} が表示されるようになっています。

YAML 形式での出力

最後に、構造体の内容を YAML 形式で出力するメソッドが生成できるようにします。cmd/protoc-gen-subgo/main.go の processMessage を次のように変更します。

func processMessage(g *protogen.GeneratedFile, m *protogen.Message) error {
  g.P("type ", m.GoIdent, " struct {")
  for _, f := range m.Fields {
    ty, err := goTypeName(f)
    if err != nil {
      return err
    }
    g.P(f.GoName, " ", ty)
  }
  g.P("}")

  // 次を追加
  g.P("func (m ", m.GoIdent, ") PrintYAML(indentSpaces int) {")
  printfIdent := protogen.GoIdent{
    GoName: "Printf",
    GoImportPath: protogen.GoImportPath("fmt"),
  }
  repIdent := protogen.GoIdent{
    GoName: "Repeat",
    GoImportPath: protogen.GoImportPath("strings"),
  }
  g.P("indent := ", repIdent, `(" ", indentSpaces)`)
  for _, f := range m.Fields {
    if f.Desc.Kind() == protoreflect.MessageKind {
      g.P(printfIdent, `("%s`, f.GoName, `: \n", indent)`)
      g.P("m.", f.GoName, ".PrintYAML(indentSpaces + 2)")
      continue
    }
    g.P(printfIdent, `("%s`, f.GoName, `: %v\n", indent, m.`, f.GoName, `)`)
  }
  g.P("}")

  return nil
}
  • GeneratedFile.P メソッドには protogen.GoIdent 型の値を与えることができます。このとき、GoImportPath に従って適切な import 文が追加されます。従ってコード生成にあたって import "fmt" 等を直接書き込む必要はありません。

再度 protoc-gen-subgo のビルドと protoc の実行を行うと、Person と Company のそれぞれに次のような PrintYAML メソッドが生成されるようになっています。引数の indentSpaces は出力される YAML のインデント(スペース数)です。

func (m Person) PrintYAML(indentSpaces int) {
  indent := strings.Repeat(" ", indentSpaces)
  fmt.Printf("%sName: %v\n", indent, m.Name)
  fmt.Printf("%sAge: %v\n", indent, m.Age)
  fmt.Printf("%sCompany: \n", indent)
  m.Company.PrintYAML(indentSpaces + 2)
}

func (m Company) PrintYAML(indentSpaces int) {
  indent := strings.Repeat(" ", indentSpaces)
  fmt.Printf("%sName: %v\n", indent, m.Name)
}

生成される person.sub.go の内容(一部抜粋)

これを ./main.go で使ってみます。

package main

import (
  proto "example.com/protoc-gen-subgo/protobuf/go"
)

func main() {
  p := proto.Person{
    Name: "Taro",
    Age:  30,
    Company: proto.Company{
      Name: "FooBarCorp",
    },
  }

  p.PrintYAML(0)
}

./main.go の内容

go run main.go の実行結果は次のようになります。

Name: Taro
Age: 30
Company:
  Name: FooBarCorp

おわりに

ソースコード全体は https://github.com/jajimajp/protoc-plugin-tutorial で公開しています。