context.Contextを自動挿入する静的解析ツール ctxfmt を書いた
これは Go言語 Advent Calendar 2023 15日目の記事です。
はじめに
株式会社CastingONEでソフトウェアエンジニアをしている @takashabe です。普段はHR領域のSaaSをGoで書いています。
みなさん context
パッケージは活用されているでしょうか。多くのパッケージがそうであるように、現在新規に書かれる関数やメソッドは context.Context
を受け取るような設計になっていることが大半かと思います。
しかし太古に書かれたコードであったり、何らかの理由によりcontextを使っていないコードも存在し、それをメンテするときには自前でcontext対応する必要があります。ある程度はgrepなどで機械的に書き換える事ができますが、context対応しているメソッドとそうでないメソッドが混在していたり、どのメソッド呼び出し元にcontextを追加すれば良いのかなど、意外と詰まるポイントが多いです。
今回はそういったcontext未対応のコードベースに対して、context.Contextを自動で挿入するためのツールである ctxfmt を書いてみたので紹介したいと思います。
Features
ユースケースとして大きく2つ考えています。まずはメソッド定義への引数追加( def
サブコマンド)、もう1つはメソッド呼び出し時の引数追加( call
サブコマンド)です。
$ ctxfmt -h
NAME:
ctxfmt - context.Context formatter
USAGE:
ctxfmt [global options] command [command options] [arguments...]
COMMANDS:
def format method definition
call format method call
help, h Shows a list of commands or help for one command
GLOBAL OPTIONS:
--help, -h show help
メソッド定義への引数追加
$ ctxfmt def -h
NAME:
ctxfmt def - format method definition
USAGE:
ctxfmt def [command options] target file or directory
OPTIONS:
--dryrun, -n dryrun (default: false)
--config value, -c value config file path
--help, -h show help
インタフェースのメソッド宣言や、メソッド定義に対して ctx context.Context
を追加することが出来ます。
例えば以下の様なコードがあり、
package main
type Interface interface {
Foo(id int)
}
type impl struct{}
func (i *impl) Foo(id int) {}
func main() {
i := &impl{}
i.Foo(1)
}
ここで ctxfmt
コマンドを実行すると
$ ctxfmt def $GOPATH/src/github.com/takashabe/ctxfmt/examples/**
以下のようなコードが生成されます。
package main
import "context"
type Interface interface {
Foo(ctx context.Context, id int)
}
type impl struct{}
func (i *impl) Foo(ctx context.Context, id int) {}
func main() {
i := &impl{}
i.Foo(1)
}
差分は見ての通りですが、 メソッド定義に ctx context.Context
が追加され、import文も追加されています。
@@ -1,12 +1,14 @@
package main
+import "context"
+
type Interface interface {
- Foo(id int)
+ Foo(ctx context.Context, id int)
}
type impl struct{}
-func (i *impl) Foo(id int) {}
+func (i *impl) Foo(ctx context.Context, id int) {}
func main() {
i := &impl{}
メソッド呼び出し時の引数追加
$ ctxfmt call -h
NAME:
ctxfmt call - format method call
USAGE:
ctxfmt call [command options] target directory
OPTIONS:
--dryrun, -n dryrun (default: false)
--pkg value, -p value package name
--config value, -c value config file path
--help, -h show help
ここまででメソッド定義にcontextを追加することは出来ました。しかしメソッド定義側の変更だけでは、呼び出し元で引数不足によるコンパイルエラーが発生してしまうでしょう。
先ほどのファイルを例にすると以下のコードの状態です。
package main
import "context"
type Interface interface {
Foo(ctx context.Context, id int)
}
type impl struct{}
func (i *impl) Foo(ctx context.Context, id int) {}
func main() {
i := &impl{}
i.Foo(1)
}
このファイルに対して go vet
してみると以下のようなエラーが出ることがわかります。
$ go vet .
# github.com/takashabe/ctxfmt/examples
vet: ./main.go:15:9: not enough arguments in call to i.Foo
have (number)
want (context.Context, int)
ということで ctxfmt
が用意している call
サブコマンドを叩くことにより、呼び出し元に context.TODO()
を追加することが出来ます。
$ ctxfmt call --pkg='github.com/takashabe/ctxfmt/examples' $GOPATH/src/github.com/takashabe/ctxfmt/examples
processed /Users/takashabe/dev/src/github.com/takashabe/ctxfmt/examples/main.go
@@ -12,5 +12,5 @@ func (i *impl) Foo(ctx context.Context, id int) {}
func main() {
i := &impl{}
- i.Foo(1)
+ i.Foo(context.TODO(), 1)
}
configによる解析対象の調整
このように任意のメソッド宣言や呼び出しにおいて、自動的にcontextを追加することが出来る様子を見てきました。
ここで勘の良いGopherは標準パッケージにあるメソッド、例えば MarshalJSON()
や String()
を実装している場合に困りそう、と考えるかもしれません。
ctxfmt
では引数が全くないメソッド宣言や、configで明示的に指定したメソッドについては処理の対象外としています。
configは例えば以下のようなyamlファイルです。
ignore_funcs:
- Scan
- Validate
これを ctxfmt
の --config
オプションに食わせることで指定のメソッドは対象外とすることが出来ます。
$ ctxfmt def --config=config.yaml $GOPATH/src/github.com/takashabe/ctxfmt/examples/**
dryrun
こういうツールではありがちですが、 --dryrun
オプションも用意しています。dryrunモードではファイルの編集を行わないことに加え、処理対象メソッドが何であるかを明示的に出力しています。
そして def
サブコマンドと --dryrun
オプションを組み合わせることで、新たに書かれたメソッドにcontextが書かれているかどうかをチェックするlinterとしても使えるかもしれません。
$ ctxfmt def --dryrun $GOPATH/src/github.com/takashabe/ctxfmt/examples/**
/Users/takashabe/dev/src/github.com/takashabe/ctxfmt/examples/main.go at line 4: Interface.Foo()
/Users/takashabe/dev/src/github.com/takashabe/ctxfmt/examples/main.go at line 9: i.Foo()
実装
最後に簡単に実装の話をしたいと思います。
まずメソッド定義をいじる def
サブコマンドはASTを走査して、インタフェースのメソッド宣言やメソッド定義で第一引数に context.Context
がないメソッドを探し、あれば x/tools/astutil
で引数を追加しています。
astutil.AddImport(fs, file, "context")
astutil.Apply(file, func(cr *astutil.Cursor) bool {
switch decl := cr.Node().(type) {
case *ast.FuncDecl:
if decl.Name != nil {
if !isIgnoreFunc(decl.Name.Name) {
if decl.Recv != nil {
if decl.Type.Params != nil && len(decl.Type.Params.List) > 0 {
if unicode.IsUpper(rune(decl.Name.Name[0])) {
if !hasContextParam(decl.Type.Params.List) {
contextParam := &ast.Field{
Names: []*ast.Ident{ast.NewIdent("ctx")},
Type: &ast.SelectorExpr{
X: ast.NewIdent("context"),
Sel: ast.NewIdent("Context"),
},
}
decl.Type.Params.List = append([]*ast.Field{contextParam}, decl.Type.Params.List...)
}
}
}
}
}
}
case *ast.TypeSpec:
if interfaceType, ok := decl.Type.(*ast.InterfaceType); ok {
for _, m := range interfaceType.Methods.List {
if method, ok := m.Type.(*ast.FuncType); ok {
if !hasContextParam(method.Params.List) {
contextParam := &ast.Field{
Names: []*ast.Ident{ast.NewIdent("ctx")},
Type: &ast.SelectorExpr{
X: ast.NewIdent("context"),
Sel: ast.NewIdent("Context"),
},
}
method.Params.List = append([]*ast.Field{contextParam}, method.Params.List...)
}
}
}
}
}
return true
}, nil)
次にメソッド呼び出しをいじる call
サブコマンドですが、こちらは力技を使っており、コンパイルエラーが出る状態を前提とし、引数不足エラーが報告されたメソッド名を編集するロジックとなっています。
本来であれば呼び出し先の型を調べて、厳密に編集をしたほうが良いですが、型情報の解析をするために依存関係の解決がややこしかったりするので妥協した形です。
以下のように packages.Load
でコンパイルのエラー情報を取得し、そのメッセージに対して正規表現でメソッド名を抜き出しています。その後はいつものとおりASTをいじる感じです。
pkgs, err := packages.Load(cfg, pkgName)
if err != nil {
return err
}
for _, pkg := range pkgs {
for _, err := range pkg.Errors {
funcNames, ok := notEnoughContextArgs(err.Error())
if !ok {
continue
}
...
var functionCallRegex = regexp.MustCompile(`not enough arguments in call to [\w.]+\b\.(\w+)`)
func notEnoughContextArgs(errMessage string) ([]string, bool) {
var funcNames []string
matches := functionCallRegex.FindAllStringSubmatch(errMessage, -1)
for _, match := range matches {
if len(match) >= 2 {
funcNames = append(funcNames, match[1])
}
}
return funcNames, len(funcNames) > 0
}
今後の展望
ひとまず作ってはみたものの、まだCIなど全く整備されていないのと、コマンド実行時の余分な引数であったり、考慮パターンも少ないので、もうちょっと使いやすいものを目指して絶賛開発中です。
特に柔軟なlinterとして価値を感じており、まずは手元でフィードバックサイクルを回していきたいと思っています。
まとめ
メソッド定義とメソッド呼び出しの双方に context.Context
を自動挿入するためのツールである ctxfmt
の紹介を行いました。
レガシーコードのリファクタや、チーム開発でのlinterとして活用していけると考えています。
いつもの
株式会社CastingONE では、ASTを眺めるのが好きなソフトウェアエンジニアを募集しております。Twitter(X)でもBlueskyでも良いのでお気軽にご連絡ください!
Discussion