Refactoring Go code using AST replacement

For complex refactoring of Go code you can use the Go AST parser and printer to find patterns and replace them with your new code, retaining parameters and variables.

A code base with a lot of simple Go tests wasn’t quite readable and we wanted to remove repetitive assertions using testify. The goal was to implement such conversions:

// before:
err := ThisMightReturnError()
if err != nil {
t.Fatalf("my error: %s", err)
}
// after:
err := ThisMightReturnError()
require.NotNil(t, "my error: %s", err)

As regexes could solve this, there is an even more interesting solution for (over)engineers: AST rewriting.

Example of an AST

Parsing the AST

Using go/parser we can parse a Go source file and get the AST from it. Spew can be used to dump the AST in a human readable format:

import (
"go/parser"
"go/token"
"github.com/davecgh/go-spew/spew"
)
fs := token.NewFileSet()
f, err := parser.ParseFile(fs, "example/main_test.go", nil, parser.AllErrors)
spew.Dump(f)

Therefore we get the AST tree that looks similar to this part:

(*ast.File)(0xc00014c000)({
Doc: (*ast.CommentGroup)(<nil>),
Package: (token.Pos) 1,
Name: (*ast.Ident)(0xc0001280c0)(main),
Decls: ([]ast.Decl) (len=3 cap=4) {
(*ast.GenDecl)(0xc000108080)({
Doc: (*ast.CommentGroup)(<nil>),
TokPos: (token.Pos) 15,
Tok: (token.Token) import,
Lparen: (token.Pos) 22,
Specs: ([]ast.Spec) (len=2 cap=2) {
(*ast.ImportSpec)(0xc00010c300)({
Doc: (*ast.CommentGroup)(<nil>),
...

Printing the code again

If we want to print the AST back as code we can use go/printer:

import "os"
import "go/printer"
...printer.Fprint(os.Stdout, fs, f)

Finding the if statement

To target if err != nil statements, we need to create a function that detects them in an AST tree. Such a statement looks like this in AST:

(*ast.IfStmt)(0xc0000241c0)({
If: (token.Pos) 80,
Init: (ast.Stmt) <nil>,
Cond: (*ast.BinaryExpr)(0xc000074570)({
X: (*ast.Ident)(0xc00005e220)(err), // left-hand-side operator
OpPos: (token.Pos) 87,
Op: (token.Token) !=, // token.NEQ
Y: (*ast.Ident)(0xc00005e240)(nil) // right-hand-side operator
})

A matching function can be implemented as follows: It gets an AST node and checks if the node is an ast.IfStmt , goes deeper into the statement and validates the operator and variable names:

func isIfErrBlock(n ast.Node) bool {
// is: if (err != nil)
if ifStmt, ok := n.(*ast.IfStmt); ok {
// is: (err != nil)
if binExpr, ok := ifStmt.Cond.(*ast.BinaryExpr); ok {
// is: !=
if binExpr.Op != token.NEQ {
return false
}
// Check left hand identifier (err)
if ident, ok := binExpr.X.(*ast.Ident); ok {
if ident.Obj == nil {
return false
}
if ident.Obj.Kind != ast.Var || ident.Name != "err" {
return false
}
}
// Check right hand identifier (nil)
if ident, ok := binExpr.Y.(*ast.Ident); ok {
if ident.Obj != nil || ident.Name != "nil" {
return false
}
}
return true
}
}
return false
}

Validate the If-Body

Additionally we validate that the if-body has only one expression inside and calls a Fatalf function. If we get this, we preserve the Fatalf arguments as we want to pass them to testify:

func isErrBody(n ast.Node) (bool, []ast.Expr) {
// check if it's really an if statement
ifStmt, ok := n.(*ast.IfStmt)
if !ok {
return false, nil
}
// check if it has only one item in the block
if len(ifStmt.Body.List) != 1 {
return false, nil
}
stmt := ifStmt.Body.List[0] // Cast expression to get the call
expr, ok := stmt.(*ast.ExprStmt)
if !ok {
return false, nil
}
call, ok := expr.X.(*ast.CallExpr)
if !ok {
return false, nil
}
// Get the function
fun, ok := call.Fun.(*ast.SelectorExpr)
if !ok {
return false, nil
}
if fun.Sel.Name != "Fatalf" {
return false, nil
}
return true, call.Args
}

Putting it all together

We can now:

  • Read and parse the AST of a Go file
  • Detect the if-statement
  • Validate the if-body, if it meets our criteria and extract the Fatalf call arguments
  • Serialize the AST as a Go file

To put it in a working piece together we use the astutil to conveniently walk through our AST and find-and-replace our if-blocks:

import "golang.org/x/tools/go/ast/astutil"...
spew.Dump(f)
astutil.Apply(f, func(cr *astutil.Cursor) bool {
// Check if it is: if err != nil
if !isIfErrBlock(cr.Node()) {
return true
}
// Check if body contains only one statement and get args for it
sigRet, args := isErrBody(cr.Node())
if !sigRet {
return true
}
// Replace values
cr.Replace(&ast.ExprStmt{
X: &ast.CallExpr{
Fun: ast.NewIdent("require.NotNil"),
Args: append([]ast.Expr{ast.NewIdent("t")}, args...),
},
})
return false
}, nil)
// Print result
printer.Fprint(os.Stdout, fs, f)

As you can see, we use our isIfErrBlock and isErrBlock functions from above to check if our node matches the criteria and if so, we replace it with an expression statement:

&ast.ExprStmt{
X: &ast.CallExpr{
Fun: ast.NewIdent("require.NotNil"),
Args: append([]ast.Expr{ast.NewIdent("t")}, args...),
},
})

which equals to require.NotNil(t, args) , where the args are taken from the original t.Fatalf(args)statement.

Example of operation

The demonstration code from above converts files like this:

package mainimport (
"errors"
"testing"
)
func ThisMightReturnError() error {
return errors.New("I'm an error")
}
func TestMain(t *testing.T) {
err := ThisMightReturnError()
if err != nil {
t.Fatalf("my error: %s", err)
}
}

to this:

package mainimport (
"errors"
"testing"
)
func ThisMightReturnError() error {
return errors.New("I'm an error")
}
func TestMain(t *testing.T) {
err := ThisMightReturnError()
require.NotNil(t, "my error: %s", err)
}

Summary

We could easily replace code blocks with other code blocks using AST parsing and serializing. The code can be improved for better reliability and usability (like importing github.com/stretchr/testify/require automatically or passing *testing.T too).

The demo code is available here: https://gist.github.com/djboris9/f58c7bf19472dab829a858fd426685e9

Linux, Go, Container, PostgreSQL, Bitcoin, Infosec

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store