Using go/analysis to fix your source code

The go/analysis library allows us to write custom Go linters in an easy and fast way. Did you know that you can also rewrite source code with it?

Using go/analysis to fix your source code

This is a followup blog post to my initial blog post about the go/analysis framework. If you haven't read the blog post, I recommend reading that first as this builds on top of it. We're going to keep it short because once you have already a linter, you can easily add the rewrite functionality.

As a recap, the go/analysis package is a Go library that allows you to write linters and checkers for Go source code's in a fast and standardized way. Usually, people use it to implement linters, which means it shows warnings and errors to you. Running it has no side effects.

However, go/analysis is also able to rewrite your source code. How does it work?

Example linter

First, let's write a simple linter that reports the usages of the + operator
in a binary expression (this is a simplified version based on the linter we
wrote in the first blog post):

package main

import (
	"github.com/fatih/addlint/addcheck"
	"golang.org/x/tools/go/analysis/singlechecker"
)

func main() {
	singlechecker.Main(addcheck.Analyzer)
}
package addcheck

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/printer"
	"go/token"

	"golang.org/x/tools/go/analysis"
)

var Analyzer = &analysis.Analyzer{
	Name: "addlint",
	Doc:  "reports integer additions",
	Run:  run,
}

func run(pass *analysis.Pass) (interface{}, error) {
	for _, file := range pass.Files {
		ast.Inspect(file, func(n ast.Node) bool {
			be, ok := n.(*ast.BinaryExpr)
			if !ok {
				return true
			}

			if be.Op != token.ADD {
				return true
			}

			pass.Report(analysis.Diagnostic{
				Pos:     be.Pos(),
				Message: "integer addition found",
			})
			return true
		})
	}

	return nil, nil
}

Let's run the linter on the this example:

$ cat test.go
package main

import (
        "fmt"
)

func main() {
        result := 3 + 2
        fmt.Println("Result", result)
}
$ go install cmd/addlint/main.go
$ addlint test.go
/Users/fatih/test.go:8:12: integer addition found

This linter is simpler than what we had in our previous blog post. For example,
it doesn't check whether the right and left hand side of the binary expressions
are integer types. But that's not important for us.

The important detail we need to know about the go/analysis package is, how
does the code rewrite work? It does it by providing two things:

  1. the start/end range of the lines we want the change
  2. the actual text we want to replace

This means that you specify the range (via token.pos) in your file and then
provide the text in []byte type to go/analysis. You don't modify the
AST! You'll find out that you still have to do it in this blog post, but that's
temporary to render the AST and is not used by the go/analysis package.


Side note: Which is better? I've worked with both approaches in the past,
rewriting the AST is easier, however, it's also not very scalable if you have
multiple fixer/linters. As an example, if you have multiple check logic in your
linter (or use a multichecker), if you do a AST rewrite in the first check, any
following checks will now use an AST that is modified and is not the original
copy anymore.  One way to mitigate is to deep-copy the AST and pass it to
multiple checkers, but that has its issues. The Go AST also has issues
retaining the comments correctly if you modify an AST, which makes it hard
to use the library.

In the text-based rewrite, all your checkers can act on an un-modified
AST and can provide the rewrite separately. Even if there are conflicts, the
user can easily select one and apply the fix that is important for them.

Note that Go's AST package is good at traversing the tree, but it doesn't have a
proper API for modifying the AST safely and reliably. There are some
nice projects, such as dst which provide a
way better way to manipulate the AST of a Go source code.


Finally we need one last thing: a helper function to render (pretty print) AST
expressions to Go source code:

// render returns the pretty-print of the given node
func render(fset *token.FileSet, x interface{}) string {
	var buf bytes.Buffer
	if err := printer.Fprint(&buf, fset, x); err != nil {
		panic(err)
	}
	return buf.String()
}

We used this in our first blog post to pretty-print binary expressions in
human-readable form, i.e: 3 + 2. This will be used to still output human
readable linter messages, but also to render the new text (source code) we want
to rewrite.

Adding a suggested fix

For our example let's rewrite all + operators with * (token.MUL)
operator.

We're going to first store the un-modified version of the binary expression,
and then change the code and finally create the modified version:

oldExpr := render(pass.Fset, be) // 3 + 2

// let's create a new text by simply changing the AST and rendering it back to
// a string. 
be.Op = token.MUL
newExpr := render(pass.Fset, be) // 3 * 2

// We could also replace the `+` character directly in the string
// without modifying the AST. Pick one that fits your case.
newExpr := strings.ReplaceAll(oldExpr, "+", "*")

The oldExpr will be used to show to the user what is wrong. newExpr will be
used to rewrite (replace) the old expression. As you see we can create the
newExpr in two different ways.

Once we have this, the only thing left is to tell go/analysis which lines we
want to change. The way we do is, we're going to set the SuggestedFixed field
of the analysis.Diagnostic type (simplified):

// A Diagnostic is a message associated with a source location or range.
type Diagnostic struct {
	// SuggestedFixes contains suggested fixes for a diagnostic which can
	// be used to perform edits to a file that address the diagnostic.
	// Diagnostics should not contain SuggestedFixes that overlap.
	// Experimental: This API is experimental and may change in the future.
	SuggestedFixes []SuggestedFix // optional
}

// A SuggestedFix is a code change associated with a Diagnostic that a user can choose
// to apply to their code. Usually the SuggestedFix is meant to fix the issue flagged
// by the diagnostic.
// TextEdits for a SuggestedFix should not overlap. TextEdits for a SuggestedFix
// should not contain edits for other packages.
// Experimental: This API is experimental and may change in the future.
type SuggestedFix struct {
	// A description for this suggested fix to be shown to a user deciding
	// whether to accept it.
	Message   string
	TextEdits []TextEdit
}

// A TextEdit represents the replacement of the code between Pos and End with the new text.
// Each TextEdit should apply to a single file. End should not be earlier in
// the file than Pos.
// Experimental: This API is experimental and may change in the future.
type TextEdit struct {
	// For a pure insertion, End can either be set to Pos or token.NoPos.
	Pos     token.Pos
	End     token.Pos
	NewText []byte
}

A couple of things here:

  • The SuggestedFix.Message field is not implemented yet. Populating it won't prompt the user and ask them anything, but we're still going to add it so our linter is future-proof.
  • Second, you can add multiple fixes for a single Diagnostic. However, that requires storing some state internally (while traversing the AST) and it's more work than needed. I've found that providing a single TextEdit is easier to work with (this is what errwrap is also using)

Now, this is how we're going to add a suggested fix:

  • Add the SuggestedFixes field to the analysis.Diagnostic struct.
  • As I said we're setting SuggestedFix.Message, even though it's not used by go/analysis yet.
  • And our SuggestedFix.NewText is set to the newly rendered expression (which contains the binary expression with the * operator). This is the newExpr variable we created above with the render() function:
pass.Report(analysis.Diagnostic{
	Pos:     be.Pos(),
	Message: fmt.Sprintf("integer addition found %q", oldExpr),
	SuggestedFixes: []analysis.SuggestedFix{
		{
			Message: fmt.Sprintf("should replace `%s` with `%s`", oldExpr, newExpr),
			TextEdits: []analysis.TextEdit{
				{
					Pos:     be.Pos(),
					End:     be.End(),
					NewText: []byte(newExpr),
				},
			},
		},
	},
})

Running our linter behaves still the same. This is great because it means that adding a SuggestedFix is backwards compatible and safe to do:

$ addlint test.go
/Users/fatih/test.go:8:12: integer addition found "3 + 2"

Now, to apply the suggested fixes, all we do is to pass the -fix flag:

$ addlint -fix test.go
/Users/fatih/test.go:8:12: integer addition found "3 + 2"

$ cat test.go
package main

import (
        "fmt"
)

func main() {
        result := 3 * 2 // <--- THIS IS CHANGED
        fmt.Println("Result", result)
}

You can verify that it works by re-running your linter again for the same file. You'll see that the linter will not report anything because it's now fixed!:

$ addlint test.go

Verdict

That's it! As you see the only change we did is to set the SuggestedFixes field. The go/analysis framework handles the rest for us (including setting the -fix flag for us).

For a real world example, check out the errwrap tool, which reports errors without the new %w verb directive, and can rewrite them if needed.