Go の関数を context 対応するツール - 詩と創作・思索のひろば

最初から完璧な設計と実装ができているなら苦労はないわけだけど、実際にはそうもいかない。具体的にはある程度の規模になってくると「あーこの関数 context.Context 対応したい！」みたいな気持ちが湧いてくるわけです。context 対応ってのは、第一引数に ctx context.Context を追加することですね。

そういうことをやるツールを書きました。

GitHub - motemen/go-ctxize: Rewrite functions to have "Context"s

go get github.com/motemen/go-ctxize/cmd/goctxize で、goctxize というバイナリが手に入ります。

サンプル

README やテストにある例だけど、

// $GOPATH/src/example.com/foo/foo.go
package foo

func F() {
}

// $GOPATH/src/example.com/foo/foo_test.go
package foo

import "testing"

func TestF(t *testing.T) {
    F()
}

// $GOPATH/src/example.com/bar/bar.go
package bar

import (
    "example.com/foo"
)

func bar() {
    foo.F()
}

という感じのコードにおいて、foo.F() のシグネチャが foo.F(ctx context.Context) であってほしい！という場合。

goctxize example.com/foo.F

としてやると foo.go が以下のように書き換わります。

// $GOPATH/src/example.com/foo/foo.go
package foo

import "context"

func F(ctx context.Context) {
}

// $GOPATH/src/example.com/foo/foo_test.go
package foo

import (
    "context"
    "testing"
)

func TestF(t *testing.T) {
    ctx := context.TODO()

    F(ctx)
}

ほんとは example.com/foo.F を呼び出しているほかのパッケージも書き換えたいわけなので、第2引数以降に追加のパッケージを指定してやると:

goctxize example.com/foo.F example.com/bar

bar.go もこうなります。

// $GOPATH/src/example.com/bar/bar.go
package bar

import (
    "context"
    "example.com/foo"
)

func bar() {
    ctx := context.TODO()

    foo.F(ctx)
}

コンテキストの推測

呼び出し側のファイルを書き換える際、*gin.Context みたいに context.Context を実装している変数がすでに呼び出し側のスコープに存在しているときはその変数を使ってくれます。かしこいですね。スコープ周辺の判定はけっこう雑。

// before
package web

import (
    "example.com/foo"
    "github.com/gin-gonic/gin"
)

func web(c *gin.Context) {
    foo.F()
}

// after
package web

import (
    "example.com/foo"
    "github.com/gin-gonic/gin"
)

func web(c *gin.Context) {
    foo.F(c)
}

また、コンテキスト化される関数の中に context.TODO() で初期化されている変数があったら、その変数を消し去ることもしてくれます。

// before
package bar

import (
        "context"

        "example.com/foo"
)

func alreadyHasCtxInside() { // これを ctx 化する
        ctx := context.TODO()
        _ = ctx
}

// after
package bar

import (
        "context"

        "example.com/foo"
)

func alreadyHasCtxInside(ctx context.Context) {

        _ = ctx
}

パッケージの指定

フルパスでパッケージを指定するのは煩雑なので、相対パスによる指定もできます:

# PWD=$GOPATH/src/example.com
goctxize ./foo.F ./bar ./baz

とか。go list と同じ解釈をするので、./... もいけます。

これは x/tools/go/packages パッケージによって実現されていて、作ってる途中で浦島太郎的にこの存在を知ったんだけど（Go 1.11 からの登場）、これがすげー便利。これまで静的解析をおこなうツールを書くとき、複数のパッケージにまたがって型情報を読み込みたいときには x/tools/go/loader を使うのが定石だった（と思う）のだけど、これは ./... 的な引数をうまく解釈することができない不都合があって、やや物足りなかったのです。Go モジュールに対応していることもあり、今後はこれを使うのが標準となる予定みたいですね。

どうぞご利用ください。