Protocol Buffers の oneof をリストする (+テンプレートで) プラグイン作った

Protocol Bufferにて、あるoneofフィールドについてケースごとに処理の振り分けをしたいことがあった。 これ、いくら探しても簡単にできるいい方法が見つからず (探し方が悪いのかもしれない)、結局簡単なプラグインを作った。

別のプログラムのサブプログラムとして作ったんだけど、他にPythonとかでさっと使える簡単な例も見つからず、微妙に苦労したので一応 protobuf-oneof-listing-plugin としてその部分だけGithubで公開しておくことにした。

簡単な実装例程度のソースなので読めば必要なさそうだけど、軽くソース抜粋して作り方をメモしておく。

まずprotobufライブラリをインストールしておく。

$ pip install protobuf

環境によっては--userなどをつける必要がある。あとレポジトリ内のやつはテンプレートによる生成のためもうちょっと色々インストール必要。

protobufのpluginをインポートする。

from google.protobuf.compiler import plugin_pb2 as plugin

リクエストはstdinから来るようなので、それを読み取って CodeGeneratorRequest としてパースする。

import sys

data = sys.stdin.buffer.read()
request = plugin.CodeGeneratorRequest()
request.ParseFromString(data)

入力されたファイルはパースされた状態でrequest.proto_fileに入っている。配列になっていて、複数個ファイルの場合もあるので注意。

リクエストパラメータ (protoc実行時に--hoge_out=XXXXX:./out/とするときのXXXXXの部分) は request.parameter に入っている。この部分はパースとかされず : の前そのままの文字列。

生成結果は同様に CodeGeneratorResponse として生成して、stdout。

response = plugin.CodeGeneratorResponse()

# ... responseに生成結果を色々書き込む ...

output = response.SerializeToString()
sys.stdout.buffer.write(output)

このresponseにはfileフィールドがあり、そこにaddしていくことで結果となるファイルの生成ができる。

f = response.file.add() # f の型は response.File
f.name = "test-output-file.txt"
f.content = "Generation result"

エラーが起こった場合は error フィールドに書く。

response.error = "エラーです"

今回作ったプラグインの主目的であるoneofの場合の処理などは protobuf-oneof-listing-plugin レポジトリ にある protoc_oneof_listing_plugin.py が参考になるかもしれない。

その他 CodeGeneratorResponseCodeGeneratorRequest のフィールドなどは Google の plugin.proto などを参考に頑張る。

実行は基本的には protoc に実行ファイル名を与えるだけ。プラグインの例えばファイル名が test_plugin.py だったら

$ protoc test.proto --plugin="protoc-gen-test=./test_plugin.py" \
    --test_out=templates=hoge:./out/

といった感じ。

参考文献