はじめに
Goには、gofmtというコマンドがあることはご存知かと思います。
そして、そのフォーマットがGoコミュニティの標準となっています。
go/formatパッケージでは、gofmtと同じスタイルのフォーマットでソースコードを整形することができます。
また、go/formatパッケージを使うと、ASTについても整形することがきます。
この記事では、ASTのノードをいじって、整形する方法について説明します。
なお、この記事を書いた2016年12月現在のGoのバージョンは1.7.4が最新です。
ASTをいじる
ASTをいじると書きましたが、具体的に何をするんでしょうか？
ここでは、AST上のノードを別のノードに入れ替えたり、ノードのフィールドを差し替えることを指しています。
たとえば、x+yという式があった場合、このxとyをそれぞれ10と20に置き換えて、10+20のような式に変えたいとします。
x+yは以下のように構成されるASTになるはずなので、これの*ast.Identの部分を入れ替えてやればよいでしょう。
*ast.BinaryExpr (+)
├── *ast.Ident (x)
└── *ast.Ident (y)
ast.BinaryExpr構造体は2項演算式を表し、オペランドはXフィールドとYフィールドにast.Expr型の値として保持されます。
つまり、XフィールドとYフィールドを差し替えて、以下のような構成のASTにしてやればよいでしょう。
*ast.BinaryExpr (+)
├── *ast.BasicLit (10)
└── *ast.BasicLit (20)
ast.BasicLit構造体は、以下のようなフィールドを持つ構造体です。
type BasicLit struct {
ValuePos token.Pos // literal position
Kind token.Token // token.INT, token.FLOAT, token.IMAG, token.CHAR, or token.STRING
Value string // literal string; e.g. 42, 0x7f, 3.14, 1e-9, 2.4i, 'a', '\x7f', "foo" or `\m\n\o`
}
そのため、Kindフィールドをtoken.INTにして、Valueフィールドを"10"と"20"にしたものをオペランドとして差し替えてやればOKです。
この時、ValuePosフィールドについては、ゼロ値、つまり0としておいて構いません。なお、値が0のtoken.Posはtoken.NoPosという名前の定数として定義されています。
上記の処理をコードに落とすと以下のようになります。
fset := token.NewFileSet()
expr, err := parser.ParseExprFrom(fset, "sample.go", `x+y`, 0)
if err != nil {
log.Fatalln("Error:", err)
}
fmt.Println("==== BEFORE ====")
ast.Print(fset, expr)
binaryExpr := expr.(*ast.BinaryExpr)
binaryExpr.X = &ast.BasicLit{
Kind: token.INT,
Value: "10",
}
binaryExpr.Y = &ast.BasicLit{
Kind: token.INT,
Value: "20",
}
fmt.Println("==== AFTER ====")
ast.Print(fset, expr)
そして、出力結果は以下のようになります。
==== BEFORE ====
0 *ast.BinaryExpr {
1 . X: *ast.Ident {
2 . . NamePos: sample.go:1:1
3 . . Name: "x"
4 . . Obj: *ast.Object {
5 . . . Kind: bad
6 . . . Name: ""
7 . . }
8 . }
9 . OpPos: sample.go:1:2
10 . Op: +
11 . Y: *ast.Ident {
12 . . NamePos: sample.go:1:3
13 . . Name: "y"
14 . . Obj: *(obj @ 4)
15 . }
16 }
==== AFTER ====
0 *ast.BinaryExpr {
1 . X: *ast.BasicLit {
2 . . ValuePos: -
3 . . Kind: INT
4 . . Value: "10"
5 . }
6 . OpPos: sample.go:1:2
7 . Op: +
8 . Y: *ast.BasicLit {
9 . . ValuePos: -
10 . . Kind: INT
11 . . Value: "20"
12 . }
13 }
ここではast.Print関数を使ってASTの構造を出力しています。
うまくオペランドが入れ替わっていることが分かるでしょう。
さて、次に変更したASTから文字列として式を得たいと思います。
フォーマットをかける
go/formatパッケージには、以下の2つの関数が用意されています。
func Node(dst io.Writer, fset *token.FileSet, node interface{}) error
func Source(src []byte) ([]byte, error)
format.Node関数は、ASTのノードに対してフォーマットをかけ、format.Source関数は、文字列で表されたコードに対してフォーマットを掛けます。
ここでは、ASTのノードを対象としているのでformat.Node関数を用います。
format.Node関数の第1引数は、出力先のio.Writeです。文字列としてコードが欲しい場合は、bytes.Buffer型などを渡すと良いでしょう。
第2引数のfsetは、パースする際に渡したtoken.FileSet型の値です。詳しくは、「ASTを取得する方法を調べる」という記事に書いてますので、そちらを御覧ください。
第3引数は、整形するASTのノードです。
さて、前置きが長くなりましたが、ast.Node関数を呼び出して、ソースコードを整形してみましょう。
先程オペランドを置き換えた、2項演算式のASTを整形してみます。
if err := format.Node(os.Stdout, fset, expr); err != nil {
log.Fatalln("Error:", err)
}
出力結果は以下のようになります。
10 + 20
うまく、xとyが10と20に入れ変わってますね。
また、ついでに+の間にスペースが入っているので整形されていることが分かります。
最後にすべてのコードを載せておきます。The Go Playgroundでも実行可能です。
package main
import (
"fmt"
"go/ast"
"go/format"
"go/parser"
"go/token"
"log"
"os"
)
func main() {
fset := token.NewFileSet()
expr, err := parser.ParseExprFrom(fset, "sample.go", `x+y`, 0)
if err != nil {
log.Fatalln("Error:", err)
}
fmt.Println("==== BEFORE ====")
ast.Print(fset, expr)
binaryExpr := expr.(*ast.BinaryExpr)
binaryExpr.X = &ast.BasicLit{
Kind: token.INT,
Value: "10",
}
binaryExpr.Y = &ast.BasicLit{
Kind: token.INT,
Value: "20",
}
fmt.Println("==== AFTER ====")
ast.Print(fset, expr)
if err := format.Node(os.Stdout, fset, expr); err != nil {
log.Fatalln("Error:", err)
}
}
おわりに
この記事では、ASTをいじる方法とその後go/formatパッケージをつかってソースコードを整形する方法について説明しました。
ソースコード中の宣言の順番をアルファベット順に並べかえたりするツールなどを作る際にきっと使えると思いますので、ぜひやってみてください。