go/analysisのSuggestedFixでコードを修正する

Goの既存コードを修正するツールを作る時、

既存コードをどう書き換えて
出力して
テストするか

を考えなければいけないのが少し面倒だと思っていました。
が、golang.org/x/tools/go/analysisのSuggestedFixを使えばすごく簡単にできてしまいます。

golang.org/x/tools/go/analysisは staticcheckやgolangci-lintなどの静的解析ツールでよく使われているパッケージです。

例えば以下のような、関数の引数にcontext.Contextがあるかどうかチェックするツールがあったとして、

func run(pass *analysis.Pass) (interface{}, error) {
    inspect := pass.ResultOf[inspect.Analyzer].(*inspector.Inspector)

    nodeFilter := []ast.Node{
        (*ast.FuncDecl)(nil),
    }

    inspect.Preorder(nodeFilter, func(n ast.Node) {
        decl := n.(*ast.FuncDecl)
        if decl.Type.Params.NumFields() > 0 {
            // NOTE: 第1引数のみを文字列でチェックしているので厳密ではない
            if types.ExprString(decl.Type.Params.List[0].Type) == "context.Context" {
                return
            }
        }

        pass.Reportf(decl.Pos(), "missing ctx in parameter")
    })

    return nil, nil
}

これを、もしチェックに引っ掛かったら引数にcontext.Contextを追加できるように変更してみます。

まずはpass.Reportfをpass.Reportに変更し、直接Diagnosticを渡せる形にします。

pass.Report(analysis.Diagnostic{
    Pos:     decl.Pos(),
    Message: "missing context in parameter",
})

そしてSuggestedFixesとしてコードを変更する場所(PosからEnd)と書き換え後のコード(NewText)を渡します。

pass.Report(analysis.Diagnostic{
    Pos:     decl.Pos(),
    Message: "missing context in parameter",
    SuggestedFixes: []analysis.SuggestedFix{{
        Message: "add ctx to parameter",
        TextEdits: []analysis.TextEdit{{
            Pos:     decl.Pos(),
            End:     decl.Type.Params.Closing + 1,
            NewText: b,
        }},
    }},
})

書き換え後のコードは標準パッケージのformat.Nodeを使って作ります。

func newText(pass *analysis.Pass, decl *ast.FuncDecl) ([]byte, error) {
    // Godoc、戻り値、関数の中身は使わずにコードを整形する
    f := &ast.FuncDecl{
        Recv: decl.Recv,
        Name: decl.Name,
        Type: &ast.FuncType{
            Params: &ast.FieldList{
                List: append([]*ast.Field{{
                    Names: []*ast.Ident{{Name: "ctx"}},
                    Type: &ast.SelectorExpr{
                        X:   &ast.Ident{Name: "context"},
                        Sel: &ast.Ident{Name: "Context"},
                    },
                }}, decl.Type.Params.List...),
            },
        },
    }

    var buf bytes.Buffer
    if err := format.Node(&buf, pass.Fset, f); err != nil {
        return nil, err
    }
    return buf.Bytes(), nil
}

この書き換えを実際に適用するにはコマンドラインツールとして実行する時に-fixフラグを付けるようにすればOKです。
なお、-fixフラグはunitcheckerだと渡せないため、main.goはsinglecheckerかmulticheckerを使う必要があります。

もしくは、goplsにAnalyzerとして組み込むことでエディタと連携して使うことも可能です。
多少作り込みが甘くても、リファクタリングする時だけ以下に追加し、go installして使ってみても良いかもしれません。 https://github.com/golang/tools/blob/gopls/v0.6.4/internal/lsp/source/options.go#L1108-L1150

vim + vim-lspは該当箇所で:LspCodeActionを実行すると呼び出せます。