Python で protobuf スキーマを直接読み込む

Jupyter から自作の gRPC サーバにリクエストを投げたいが、そのためには protobuf スキーマを protoc でコンパイルして生成された Python コードを import しなければならない。

例えばこんな感じで。

$ pip install grpcio grpcio-tools
$ python -m grpc_tools.protoc -I./ --python_out=. --grpc_python_out=. foo.proto
    → foo_pb2.py と foo_pb2_grpc.py が生成される

import grpc
import foo_pb2
import foo_pb2_grpc

普通の Jupyter サーバでの作業なら一度コンパイルしたものを使いまわせばいいが、Google Colab で作業をしていると時間が経つとランタイムがリセットされて消えてしまって毎回コンパイルし直す必要があり面倒臭い。

これを自動化できないかと考えた。

実装

grpc-tools が必要なので pip install grpcio-tools などしてインストールしておく。

import sys, os
from tempfile import TemporaryDirectory
from io import IOBase
from importlib.machinery import SourceFileLoader
from grpc_tools import _protoc_compiler

def load_proto(proto, name='tmp'):
    with TemporaryDirectory(dir='.') as dname:
        fname = name + '.proto'
        fpath = os.path.join(dname, fname)
        with open(fpath, 'w') as f:
            if isinstance(proto, IOBase):
                f.write(proto.read())
            else:
                f.write(proto)

        args = [
            '--proto_path=' + dname,
            '--python_out=.',
            '--grpc_python_out=.',
            fpath,
        ]
        res = _protoc_compiler.run_main([arg.encode() for arg in args])

        pb2 = SourceFileLoader(name + '_pb2', os.path.join(dname, name + '_pb2.py')).load_module()
        pb2_grpc = SourceFileLoader(name + '_pb2_grpc', os.path.join(dname, name + '_pb2_grpc.py')).load_module()

    return pb2, pb2_grpc

一時ディレクトリを作成して、そこに proto ファイルを置いてコンパイルして、生成された Python コードを import して返す load_proto という関数を作った。

使う

ファイルを読み込む場合

with open('foo.proto') as f:
    foo_pb2, foo_pb2_grpc = load_proto(f)

文字列を読み込む場合

proto = """
    syntax = "proto3";

    package foo;

    service Foo {
        ...
    }
    """
foo_pb2, foo_pb2_grpc = load_proto(proto)

proto ファイルを明示的にコンパイルしなくても Python で扱えるようになった。

注意点

これ普通に動くけど、あまりやってはいけないことをやっている気がするので真似する際は注意して欲しい。