[Go] `context.Context` が `ctx` という名前で、関数の第一引数に渡されているかを調べる静的解析ツールを作った

目的

• contextパッケージの公式ドキュメントに Context について以下のように記載されています

Do not store Contexts inside a struct type; instead, pass a Context explicitly to each function that needs it. The Context should be the first parameter, typically named ctx:

ざっくりとした意味としては

Context は構造体に埋め込むべきではなく、それそれの関数に明示的に渡すべきである。関数の第一引数にすべきで、ctx という名前をつけるのが典型的だ。

• 個人的にはこれが一種のコーディングルールのような感じになっているので、静的解析でチェックできるようにすれば楽そうだと思いました

実装

テストコード

まずはテストコードから書いてみました。
より静的解析っぽくするために context パッケージをあえて ctxPkg として定義してimportしたテストファイルです。
関数名から分かるように OKFunc1 と OKFunc2 は

• ctx という変数名で context.Context を定義している
• 関数の第一引数に渡されている

という点で公式ドキュメントにあるような書き方になっています。しかし、 NGFunc1 と NGFunc2 は変数名が正しくなかったり、引数の位置が正しくない位置にあります。

_example.go

package example

import (
	"context"
	ctxPkg "context"
)

func OKFunc1(ctx context.Context, id int, email string) {}

func OKFunc2(ctx ctxPkg.Context, attrs map[string]interface{}) {}

func NGFunc1(id int, email string, ctx context.Context) {}

func NGFunc2(c context.Context) {}

解析コード

大まかなロジックとしては

• ファイル全体を解析
• 解析の対象を関数の宣言だけに絞る
• 変数が context.Context のinterfaceを満たしているかチェック
• 満たしているなら、変数名と引数の位置をチェック

というような流れになります。できたコードは以下になります。

package main

import (
	"fmt"
	"go/ast"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"log"
)

func main() {
	fset := token.NewFileSet()

	file, err := parser.ParseFile(fset, "_example.go", nil, 0)
	if err != nil {
		log.Fatal("Error: ", err)
	}

	info := &types.Info{
		Defs: map[*ast.Ident]types.Object{},
	}

	config := &types.Config{
		Importer: importer.Default(),
	}

	_, err = config.Check("main", fset, []*ast.File{file}, info)
	if err != nil {
		log.Fatal("Error:", err)
	}

	p, err := config.Importer.Import("context")
	if err != nil {
		log.Fatal("Error:", err)
	}
	ctxType := p.Scope().Lookup("Context").Type()
	it, ok := ctxType.Underlying().(*types.Interface)
	if !ok {
		log.Fatal("should be found Context interface")
	}

	ast.Inspect(file, func(n ast.Node) bool {
		f, ok := n.(*ast.FuncDecl)
		if !ok {
			return true
		}

		for i, p := range f.Type.Params.List {
			ident := p.Names[0]

			obj := info.ObjectOf(ident)
			if obj == nil {
				return true
			}

			if types.Implements(obj.Type(), it) {
				if ident.Name != "ctx" {
					fmt.Printf("%s: variable name of context.Context in `%s` is invalid\n", fset.Position(ident.Pos()), f.Name)
				}
				if i != 0 {
					fmt.Printf("%s: position of context.Context in `%s` is invalid\n", fset.Position(ident.Pos()), f.Name)
				}
			}
		}
		return true
	})
}

解説

下準備

下準備として以下を行なっています

• token.NewFileSet() でファイル情報を記録する構造体を初期化
• parser.ParseFile の部分で _example.go をパースします
  • 最後の引数はパーサーのモードを指定する数字でコメントを含めてパースしたり、importの宣言部分だけをパースしたりすることができます
• types.Info は型チェックの結果が保持される構造体です
  • ここでは関数の宣言以外は不要な情報なので Defs のみを指定します
  • そのほかには例えば変数の使用された情報を保持する Uses などが指定できます
• types.Config で型チェックの設定を指定できます
• config.Check で型情報のチェックが走ります

	fset := token.NewFileSet()

	file, err := parser.ParseFile(fset, "_example.go", nil, 0)
	if err != nil {
		log.Fatal("Error: ", err)
	}

	info := &types.Info{
		Defs: map[*ast.Ident]types.Object{},
	}

	config := &types.Config{
		Importer: importer.Default(),
	}

	_, err = config.Check("main", fset, []*ast.File{file}, info)
	if err != nil {
		log.Fatal("Error:", err)
	}

context.Context のインポート

関数に渡される引数が context.Context なのかは「引数が context.Contextのinterfaceを満たしているか」で判断するので、元となる context.Context interfaceの情報を持ってこなくてはなりません。

参考になる記事がなくて苦労しましたが、他のパッケージから型情報を持ってくるには config.Importer.Import を使えば良いようです。Importer.Import の返り値は *types.Package です。この types.Package 構造体はGoのパッケージの情報を持っています。

さらに types.Package.Scope() の返り値がパッケージレベルで定義されているオブジェクト (関数など) の情報を持っているので、Context というinterfaceを Lookup メソッドを使って検索しています。

Lookup を行なった時点ではそれが何の型なのか分からないため、最後に Context の型情報を持ってきたのちに types.Interface に型アサーションしています。

	p, err := config.Importer.Import("context")
	if err != nil {
		log.Fatal("Error:", err)
	}
	ctxType := p.Scope().Lookup("Context").Type()
	it, ok := ctxType.Underlying().(*types.Interface)
	if !ok {
		log.Fatal("should be found Context interface")
	}

引数のチェック

いよいよ最後の部分です。ast.Inspect を使ってASTを深さ優先で探索していきます。

関数以外のノードは不要なため、最初に ast.FuncDecl に型アサーションして満たさないものは後続の処理を行わないようにしてます。

関数の引数は FuncDecl.Type.Params で取得できます。返り値としては FieldList という構造体のため、さらにそのフィールドの []*Field をfor loopで検査していきます。

info が型情報を持っているため、引数に対応する型情報を info.ObjectOf で取得します。返り値は types.Object となります。実際は types.Object は型情報以外にもパッケージ等の情報を持っています。

型情報だけが欲しいので obj.Type() で型情報を取得した後に types.Implements で types.Type が types.Interface を実装しているかチェックします。

その後は名前が ctx かどうか、引数の位置が最初になっているかどうかをチェックしてます。

	ast.Inspect(file, func(n ast.Node) bool {
		f, ok := n.(*ast.FuncDecl)
		if !ok {
			return true
		}

		for i, p := range f.Type.Params.List {
			ident := p.Names[0]

			obj := info.ObjectOf(ident)
			if obj == nil {
				return true
			}

			if types.Implements(obj.Type(), it) {
				if ident.Name != "ctx" {
					fmt.Printf("%s: variable name of context.Context in `%s` is invalid\n", fset.Position(ident.Pos()), f.Name)
				}
				if i != 0 {
					fmt.Printf("%s: position of context.Context in `%s` is invalid\n", fset.Position(ident.Pos()), f.Name)
				}
			}
		}
		return true
	})

コードを実行してみる

❯ go run ./main.go
_example.go:12:36: position of context.Context in `NGFunc1` is invalid
_example.go:14:14: variable name of context.Context in `NGFunc2` is invalid

ちゃんと NGFunc1 と NGFunc2 のみがエラーとして検出されています。

まとめ

• Goの静的解析用のパッケージを使用して、 context.Context が関数の第一引数に ctx として指定されているかどうかを調べるツールができました
• 結構簡単にできるので楽しかったです！！
• 次は analysis.Analyzer に組み込んでみたいと思います