Add initial fmt.Sprintf rewriting

Author: xdg

internal/rewriter/rewriter.go

@@ -31,34 +31,18 @@ func Rewrite(filename string, oldSource []byte) ([]byte, error) {
 }
 
 func visitor(n ast.Node) (ast.Node, bool) {
-	c, ok := n.(*ast.CallExpr)
-	if !ok {
+	c, name := getCallExprLiteral(n)
+	if c == nil {
 		return n, true
 	}
-
-	s, ok := c.Fun.(*ast.SelectorExpr)
-	if !ok {
-		return n, true
-	}
-
-	i, ok := s.X.(*ast.Ident)
-	if !ok {
-		return n, true
-	}
-
-	if i.Name != "errors" {
-		return n, true
-	}
-
-	switch s.Sel.Name {
-	case "Wrap":
+	switch name {
+	case "errors.Wrap":
 		return rewriteWrap(c), true
-	case "Wrapf":
+	case "errors.Wrapf":
 		return rewriteWrap(c), true
 	default:
+		return n, true
 	}
-
-	return n, true
 }
 
 func rewriteWrap(ce *ast.CallExpr) *ast.CallExpr {
@@ -67,6 +51,12 @@ func rewriteWrap(ce *ast.CallExpr) *ast.CallExpr {
 	copy(newArgs, ce.Args[1:])
 	newArgs = append(newArgs, ce.Args[0])
 
+	// If the format string is a fmt.Sprintf call, we can unwrap it.
+	c, name := getCallExprLiteral(newArgs[0])
+	if c != nil && name == "fmt.Sprintf" {
+		newArgs = append(c.Args, newArgs[1:]...)
+	}
+
 	// 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:
@@ -87,6 +77,25 @@ func rewriteWrap(ce *ast.CallExpr) *ast.CallExpr {
 	return newErrorfExpr(newArgs)
 }
 
+func getCallExprLiteral(n ast.Node) (*ast.CallExpr, string) {
+	c, ok := n.(*ast.CallExpr)
+	if !ok {
+		return nil, ""
+	}
+
+	s, ok := c.Fun.(*ast.SelectorExpr)
+	if !ok {
+		return nil, ""
+	}
+
+	i, ok := s.X.(*ast.Ident)
+	if !ok {
+		return nil, ""
+	}
+
+	return c, i.Name + "." + s.Sel.Name
+}
+
 func newErrorfExpr(args []ast.Expr) *ast.CallExpr {
 	return &ast.CallExpr{
 		Fun: &ast.SelectorExpr{

internal/rewriter/rewriter_test.go

@@ -15,9 +15,11 @@ func TestRewrite(t *testing.T) {
 		"wrap_string",
 		"wrap_fcn",
 		"wrap_var",
+		"wrap_sprintf",
 		"wrapf_string",
 		"wrapf_fcn",
 		"wrapf_var",
+		"wrapf_sprintf",
 	}
 
 	for _, c := range cases {

internal/rewriter/testdata/expect/wrap_sprintf.go

@@ -0,0 +1,12 @@
+package testdata
+
+import (
+	"fmt"
+	"log"
+)
+
+func main() {
+	err := fmt.Errorf("this is an error")
+	foo := "foo"
+	log.Print(fmt.Errorf("error occurred '%s': %w", foo, err))
+}

internal/rewriter/testdata/expect/wrapf_sprintf.go

@@ -0,0 +1,12 @@
+package testdata
+
+import (
+	"fmt"
+	"log"
+)
+
+func main() {
+	err := fmt.Errorf("this is an error")
+	foo := "foo"
+	log.Print(fmt.Errorf("error occurred '%s': %w", foo, err))
+}

internal/rewriter/testdata/input/wrap_sprintf.go

@@ -0,0 +1,14 @@
+package testdata
+
+import (
+	"fmt"
+	"log"
+
+	"github.com/pkg/errors"
+)
+
+func main() {
+	err := fmt.Errorf("this is an error")
+	foo := "foo"
+	log.Print(errors.Wrap(err, fmt.Sprintf("error occurred '%s'", foo)))
+}

internal/rewriter/testdata/input/wrapf_sprintf.go

@@ -0,0 +1,14 @@
+package testdata
+
+import (
+	"fmt"
+	"log"
+
+	"github.com/pkg/errors"
+)
+
+func main() {
+	err := fmt.Errorf("this is an error")
+	foo := "foo"
+	log.Print(errors.Wrapf(err, fmt.Sprintf("error occurred '%s'", foo)))
+}