At work, my team decided to switch our codebase from
pkg/errors to Go 1.13’s native error
wrapping. We used to wrap our errors like this:
var errBadStuff := errors.New("something happened")
...
err := errors.Wrapf(errBadStuff, "some context '%s'", label)Starting in 1.13, we can wrap errors directly using
fmt.Errorf and the %w code:
err := fmt.Errorf("some context '%s': %w", label, errBadStuff)Rather than change hundreds of wrappings across tens of thousands of lines by hand or with kludgy string replacement, I wrote a program to rewrite our source files using Go’s AST instead.
Parsing, walking and rewriting an AST 🔗︎
The central function of the program takes in old source, parses it, transforms
it and returns the new code. Parsing uses Go’s built-in packages token and
parser to produce an AST for a file:
func Rewrite(filename string, oldSource []byte) ([]byte, error) {
fset := token.NewFileSet()
oldAST, err := parser.ParseFile(fset, filename, oldSource, parser.ParseComments)
if err != nil {
return nil, fmt.Errorf("error parsing %s: %w", filename, err)
}
...
}Confusingly, the parser.ParseComments argument to ParseFile means “include
comments in the parse, not just code”, which I need to preserve all the input
for rewriting.
I used Fatih Arslan’s astrewrite
library for the AST manipulation, as it provides a Walk function that
(unlike Go’s built-in AST walker), lets me return a modified node to replace
in the AST structure.
func Rewrite(filename string, oldSource []byte) ([]byte, error) {
...
newAST := astrewrite.Walk(oldAST, visitor)
...
}Given the rewritten AST, the output is rendered using the same library used by
gofmt, called format:
func Rewrite(filename string, oldSource []byte) ([]byte, error) {
...
buf := &bytes.Buffer{}
err = format.Node(buf, fset, newAST)
if err != nil {
return nil, fmt.Errorf("error formatting new code: %w", err)
}
return buf.Bytes(), nil
}That’s the high-level flow. All the interesting details are in the visitor
function that finds the nodes of interest and rewrites them.
Inspecting the AST 🔗︎
Finding nodes in the AST and changing them requires understanding what the AST
looks like. My first attempt just dumped out nodes to stdout in visitor,
but I later found an online Go AST Viewer that
makes it easy to see the AST. By giving it a small bit of code, I can see how
that looks in AST-form.
For example, here is errors.Wrapf("foo '%s', err) in its AST form, which
appears as an *ast.CallExpr:
X: *ast.CallExpr {
Fun: *ast.SelectorExpr {
X: *ast.Ident {
NamePos: 10:2
Name: "errors"
}
Sel: *ast.Ident {
NamePos: 10:9
Name: "Wrapf"
}
}
Lparen: 10:14
Args: []ast.Expr (len = 2) {
0: *ast.BasicLit {
ValuePos: 10:15
Kind: STRING
Value: "\"foo '%s'\""
}
1: *ast.Ident {
NamePos: 10:27
Name: "err"
Obj: *(obj @ 54)
}
}
Ellipsis: -
Rparen: 10:30
}The important parts are the Fun and Args fields, which give the command
“name” and arguments, respectively. The name of the call is in a
SelectorExpr, split into X and Sel, but the AST dump is slightly
deceptive because SelectorExpr actually has an interface type for X:
type SelectorExpr struct {
X Expr // expression
Sel *Ident // field selector
}So when checking if a CallExpr.SelectorExpr.X is errors.Wrapf, I must
check if the concrete type is Ident.
Here’s a helper function I wrote to
extract the function call name out of a CallExpr:
func getCallExprLiteral(c *ast.CallExpr) string {
s, ok := c.Fun.(*ast.SelectorExpr)
if !ok {
return ""
}
i, ok := s.X.(*ast.Ident)
if !ok {
return ""
}
return i.Name + "." + s.Sel.Name
}If the CallExpr doesn’t give us a literal function name, the helper returns
an empty string to signal that no literal name was found.
Finding a node to rewrite 🔗︎
Knowing the AST structure and having a helper to extract a function name from
the AST is all that’s needed for the “find a node to rewrite” part of the
visitor function. If it’s visiting a CallExpr, the visitor just delegates
to a handler for that type.
func visitor(n ast.Node) (ast.Node, bool) {
switch v := n.(type) {
case *ast.CallExpr:
return handleCallExpr(v)
default:
return n, true
}
}The handler only has something to do if the function name matches the
names I’m looking for from pkg/errors.
func handleCallExpr(ce *ast.CallExpr) (ast.Node, bool) {
name := getCallExprLiteral(ce)
switch name {
case "errors.Wrap":
return rewriteWrap(ce), true
case "errors.Wrapf":
return rewriteWrap(ce), true
default:
return ce, true
}
}If it does match, then the program rewrites it. Otherwise, nodes are returned
unchanged from visitor.
Rewriting an AST node 🔗︎
Here are the signatures for the pkg/errors functions to rewrite:
func Wrap(err error, message string) error
func Wrapf(err error, format string, args ...interface{}) errorThree things need to happen to rewrite these to fmt.Errorf. First, the error
argument needs to rotate from the beginning to the end of the argument list.
Second, the message or format strings need to have ": %w" appended to them.
Finally, the function name has to change.
Everything needed for the first two changes is in the *ast.CallExpr.Args
field. For the rotation, I create a new argument list, copy everything but
the original first argument, then append the original first argument.
func rewriteWrap(ce *ast.CallExpr) *ast.CallExpr {
// Rotate err to the end of a new args list
newArgs := make([]ast.Expr, len(ce.Args)-1)
copy(newArgs, ce.Args[1:])
newArgs = append(newArgs, ce.Args[0])
...
}For appending ": %w", I have to account for whether the message/format is a
literal string or if it’s another expression, such a global variable or
function that returns a string.
If it’s a literal, then the literal value itself needs to be amended. Here’s
how a format string looks as an *ast.BasicLit type from the CallExpr
example above:
*ast.BasicLit {
ValuePos: 10:15
Kind: STRING
Value: "\"foo '%s'\""
}Note that the Value field includes the quotation marks. To append, I can
remove the trailing " and add : %w".
If the value isn’t a BasicLit, then given some expression X, I need to
replace that with the expression X + ": %w". In AST terms, that’s an
*ast.BinaryExpr that adds the original argument and a new BasicLit.
(How did I know that? I guessed, based on my experience with other language AST’s, but if I didn’t know that, I could put an example into the AST Viewer and see what it shows.)
Here’s how that looks in code. If the first argument is a basic literal,
amend the Value, otherwise, replace the first argument with a new
BinaryExpr:
func rewriteWrap(ce *ast.CallExpr) *ast.CallExpr {
...
// If the format string is a literal, we can rewrite it:
// "......" -> "......: %w"
// Otherwise, we replace it with a binary op to add the wrap code:
// SomeNonLiteral -> SomeNonLiteral + ": %w"
fmtStr, ok := newArgs[0].(*ast.BasicLit)
if ok {
// Strip trailing `"` and append wrap code and new trailing `"`
fmtStr.Value = fmtStr.Value[:len(fmtStr.Value)-1] + `: %w"`
} else {
binOp := &ast.BinaryExpr{
X: newArgs[0],
Op: token.ADD,
Y: &ast.BasicLit{Kind: token.STRING, Value: `": %w"`},
}
newArgs[0] = binOp
}
...
}The last part of the rewrite is to return a new CallExpr with fmt.Errof as
the function name and with the new argument list.
func rewriteWrap(ce *ast.CallExpr) *ast.CallExpr {
...
return newErrorfExpr(newArgs)
}
func newErrorfExpr(args []ast.Expr) *ast.CallExpr {
return &ast.CallExpr{
Fun: &ast.SelectorExpr{
X: &ast.Ident{Name: "fmt"},
Sel: &ast.Ident{Name: "Errorf"},
},
Args: args,
}
}That’s it! The visitor function will now rewrite errors.Wrap and
errors.Wrapf to fmt.Errorf.
Applying it to the codebase 🔗︎
The main function I wrote only handles individual files. While I didn’t
show it in this article, it also replaces the github.com/pkg/errors import
statement with errors so that errors.New uses the core errors library.
But it’s possible that errors isn’t needed anymore, only fmt, so in the
end, I used a loop to rewrap the errors with my program and then wash the
results through goimports:
for f in $(find . -iname "*.go"); do go-rewrap-errors -w $f; goimports -w $f; doneThis did exactly what I wanted – rewrite hundreds of wrapping functions across thousands of lines of code, without messing with my editor or string replacement. And it’s a reusable tool for other codebases at my company (or that you can use with the link at the end of this article).
I did have to go through and fix up a couple things by hand:
- Our custom error types had a
Causemethod to unwrap the inner error. I changed those toUnwrap, which is what built-inerrorsexpects. There were so few of these that it was faster to edit by hand than do it via AST transformation. - In a few places, we were using
pkg/errors.Causeto unwrap errors. I changed those to use built-inerrors.Asorerrors.Is.
We weren’t using other feature of pkg/errors so there was nothing else to do.
While the Go AST seems intimidating at first, it’s an extremely powerful tool for code transformation. With some visualization to see what before and after code looks like, and some trial and error, it made a hard job updating our codebase seem very easy.
The full code for this article may be found at github.com/xdg-go/go-rewrap-errors