🦁

context.Contextを自動挿入する静的解析ツール ctxfmt を書いた

2023/12/15に公開

これは Go言語 Advent Calendar 2023 15日目の記事です。

はじめに

株式会社CastingONEでソフトウェアエンジニアをしている @takashabe です。普段はHR領域のSaaSをGoで書いています。

みなさん context パッケージは活用されているでしょうか。多くのパッケージがそうであるように、現在新規に書かれる関数やメソッドは context.Context を受け取るような設計になっていることが大半かと思います。
しかし太古に書かれたコードであったり、何らかの理由によりcontextを使っていないコードも存在し、それをメンテするときには自前でcontext対応する必要があります。ある程度はgrepなどで機械的に書き換える事ができますが、context対応しているメソッドとそうでないメソッドが混在していたり、どのメソッド呼び出し元にcontextを追加すれば良いのかなど、意外と詰まるポイントが多いです。

今回はそういったcontext未対応のコードベースに対して、context.Contextを自動で挿入するためのツールである ctxfmt を書いてみたので紹介したいと思います。

https://github.com/takashabe/ctxfmt

Features

ユースケースとして大きく2つ考えています。まずはメソッド定義への引数追加( def サブコマンド)、もう1つはメソッド呼び出し時の引数追加( call サブコマンド)です。

help
$ 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

メソッド定義への引数追加

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 を追加することが出来ます。

例えば以下の様なコードがあり、

$GOPATH/src/github.com/takashabe/ctxfmt/examples/main.go
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/**

以下のようなコードが生成されます。

$GOPATH/src/github.com/takashabe/ctxfmt/examples/main.go
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{}

メソッド呼び出し時の引数追加

help
$ 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を追加することは出来ました。しかしメソッド定義側の変更だけでは、呼び出し元で引数不足によるコンパイルエラーが発生してしまうでしょう。
先ほどのファイルを例にすると以下のコードの状態です。

$GOPATH/src/github.com/takashabe/ctxfmt/examples/main.go
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ファイルです。

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 の紹介を行いました。

https://github.com/takashabe/ctxfmt

レガシーコードのリファクタや、チーム開発でのlinterとして活用していけると考えています。

いつもの

株式会社CastingONE では、ASTを眺めるのが好きなソフトウェアエンジニアを募集しております。Twitter(X)でもBlueskyでも良いのでお気軽にご連絡ください!

https://www.wantedly.com/projects/1130967
https://www.wantedly.com/projects/1063903

Discussion