特定の構造体の文字列を置き換えるFormatterを作ったときにgo/packagesを使って作ってみたので備忘録。

全体の流れ Link to heading

  1. packages.Loadを使ってパッケージ情報を読み込む
  2. 読み込んだpackages.Packageからast.File(Syntax)と型情報(TypesInfo)を取り出す
  3. ファイルごとに置き換え処理
  4. 保存

packages.Loadでパッケージ情報の読み込み Link to heading

packages.Load は指定したパターンからパッケージの情報を読み込める。

細かい仕様はドキュメントに任せるとして、ざっくりと説明をすると以下のようにGoでよくあるパッケージの指定方法の文字列を渡すといい感じにパッケージ情報を取得してくれる。

func main() {
	cfg := &packages.Config{
		Mode: packages.NeedTypes | packages.NeedSyntax | packages.NeedTypesInfo | packages.NeedFiles,
	}
	pkgs, err := packages.Load(cfg, "./...")
	// ~~~
}

packages.Packageの中にはConfigで指定したModeに沿った情報が格納されているため、これらを利用することで対象の特定や置き換えしたいファイルのastが取得できる。

余談だが、この方法で実装する前は filepath.Walk を利用していたが、比較にならないぐらいpackages.Loadのほうが早かった。

packages.PackageからSyntaxとTypesInfoの取り出し Link to heading

読み込んだpackages.PackageにはModeで指定した情報がいくつか含まれているが、特定の型を探して置き換えを行いたい場合は、SyntaxとTypesInfoが利用できる。

Syntaxにはそのパッケージ内に含まれるファイルがast.File型で格納されており、TypesInfoはそのパッケージで使用されている型の情報が go/types.Info 型として格納されている。

ast.File型は標準パッケージのものなので、そのままast.Inspectに渡してファイル構造を探索できる。

go/types.Info型は内部にast.Exprやast.Node、ast.Identをキーにして型の情報をmapで保持している。 そのため、ast.Inspectで探索中にそれらの型にキャストした上でmapのキーに渡すとそれに対応した情報を取得できる。

例えば、 cloud.google.com/go/spanner.Statement を利用している場所を特定したい場合、以下のようにすることパッケージ名と型名の完全一致で利用箇所を特定できる。

func Format(pkg *packages.Package, file *ast.File) {
	// ~~~
	ast.Inspect(file, func(n ast.Node) bool {
		compositeLit, ok := n.(*ast.CompositeLit)
		if !ok {
			return true
		}

		selectorExpr, ok := compositeLit.Type.(*ast.SelectorExpr)
		if !ok {
			return true
		}

		use, ok := pkg.TypesInfo.Uses[selectorExpr.Sel]
		if !ok {
			return true
		}

		if use.Type().String() != "cloud.google.com/go/spanner.Statement" {
			return true
		}
		// ~~~
	}
}

実行する側は以下のような感じ。

func run() {
	for _, pkg := range pkgs {
		for _, file := range pkg.Syntax {
			Format(pkg, file)
		}
	}
}

あとは普通の抽象構文木なので、よしなに値を操作すればfileの構造を変更できる。

保存 Link to heading

go/packages自体はあくまで解析の機能しか持っていないので、保存等はast.Fileの内容を文字列に変換して保存する必要がある。

とはいえそんなに難しい話でもなく、気をつける点といえばFileSetはpackages.Packageに含まれているものを利用するぐらいで、それ以外は普通に抽象構文木をファイルに保存する処理をすれば良い。

func Save(pkg *packages.Package, file *ast.File) {
	path := pkg.Fset.Position(file.Pos()).Filename
    // ~~~
	var buf bytes.Buffer
	if err := printer.Fprint(&buf, pkg.Fset, file); err != nil {
		panic(err)
	}

	result, err := format.Source(buf.Bytes())
	if err != nil {
		panic(err)
	}

	if err := os.WriteFile(path, result, 0644); err != nil {
		panic(err)
	}
}