allow replacing import path prefixes

.travis.yml

@@ -1,5 +1,5 @@
 language: go
 
 go:
-    - 1.4
+    - "1.10"
     - tip

README.md

@@ -9,23 +9,22 @@ Tool to rewrite import paths and [package import comments](https://golang.org/s/
 go get github.com/dmitris/prewrite
 
 # Usage
-prewrite -p prefix [-r] [-v] [path ...]
+prewrite -p prefix [-r] [-v] [path]
 
 # Command-line arguments
-* -p prefix -- prefix to add to imports and package import comments or remove (with -r) - required
-* -r        -- remove the given prefix from import statements and package import comments
+* -from <oldprefix> -to <newprefix> -- rewrite import paths replacing oldprefix with newprefix
 * -v        -- verbosely print the names of the changed files
 
-If not provided, the path defaults to the current directory (will recursively traverse). Multiple targets can be given.
+Example: `prewrite -from go.oldcompany.com -new go.newcompany.com`.
 
-The last target parameter can be either a single file or a directory (such as a root of a source tree).
+If not provided, the path defaults to the current directory (will recursively traverse).  Either a single file or a directory (such as a root of a source tree) can be given.
 
 # Examples
 
-Add a prefix to all imports (except the standard library) and package comment paths under the current directory:  
-prewrite -p go.corp.company.com/x -v
+Change the prefix in all the imports (except the standard library) and package comment paths under the current directory:  
+prewrite -from go.stealthy.com -to go.nextunicorn.com
 
-Remove a prefix from all imports and package comment paths under the current directory:  
-prewrite -p go.corp.company.com/x -r -v
+Remove a prefix from all imports and package comment paths under the directory /tmp/foobar :  
+prewrite -from go.stealthy.company.com/go.theunicorn.com -to go.nextunicorn.com /tmp/foobar
 
 

astmod/rewrite.go

@@ -22,25 +22,21 @@ import (
 // (The type of the argument for the src parameter must be string, []byte, or io.Reader.)
 //
 // return of nil, nil (no result, no error) means no changes are needed
-func Rewrite(fname string, src interface{}, prefix string, remove bool) (buf *bytes.Buffer, err error) {
+func Rewrite(fname string, src interface{}, from, to string) (buf *bytes.Buffer, err error) {
 	// Create the AST by parsing src.
 	fset := token.NewFileSet() // positions are relative to fset
 	f, err := parser.ParseFile(fset, fname, src, parser.ParseComments)
 	if err != nil {
 		log.Printf("Error parsing file %s, source: [%s], error: %s", fname, src, err)
 		return nil, err
 	}
-	// normalize the prefix ending with a trailing slash
-	if prefix[len(prefix)-1] != '/' {
-		prefix += "/"
-	}
 
-	changed, err := RewriteImports(f, prefix, remove)
+	changed, err := RewriteImports(f, from, to)
 	if err != nil {
 		log.Printf("Error rewriting imports in the AST: file %s - %s", fname, err)
 		return nil, err
 	}
-	changed2, err := RewriteImportComments(f, fset, prefix, remove)
+	changed2, err := RewriteImportComments(f, fset, from, to)
 	if err != nil {
 		log.Printf("Error rewriting import comments in the AST: file %s - %s", fname, err)
 		return nil, err
@@ -55,8 +51,8 @@ func Rewrite(fname string, src interface{}, prefix string, remove bool) (buf *by
 
 // RewriteImports rewrites imports in the passed AST (in-place).
 // It returns bool changed set to true if any changes were made
-// and non-nil err on error
-func RewriteImports(f *ast.File, prefix string, remove bool) (changed bool, err error) {
+// and non-nil err on error.
+func RewriteImports(f *ast.File, from, to string) (changed bool, err error) {
 	for _, impNode := range f.Imports {
 		imp, err := strconv.Unquote(impNode.Path.Value)
 		if err != nil {
@@ -67,24 +63,18 @@ func RewriteImports(f *ast.File, prefix string, remove bool) (changed bool, err
 		if !strings.Contains(imp, ".") || strings.HasPrefix(imp, ".") {
 			continue
 		}
-		if remove {
-			if strings.HasPrefix(imp, prefix) {
-				changed = true
-				impNode.Path.Value = strconv.Quote(imp[len(prefix):])
-			}
-		} else {
-			// if import does not start with the prefix already, add it
-			if !strings.HasPrefix(imp, prefix) {
-				changed = true
-				impNode.Path.Value = strconv.Quote(prefix + imp)
-			}
+
+		if strings.HasPrefix(imp, from) {
+			changed = true
+			newimp := strings.Replace(impNode.Path.Value, from, to, 1)
+			impNode.Path.Value = newimp
 		}
 	}
 	return
 }
 
 // RewriteImportComments rewrites package import comments (https://golang.org/s/go14customimport)
-func RewriteImportComments(f *ast.File, fset *token.FileSet, prefix string, remove bool) (changed bool, err error) {
+func RewriteImportComments(f *ast.File, fset *token.FileSet, from, to string) (changed bool, err error) {
 	pkgpos := fset.Position(f.Package)
 	// Print the AST.
 	// ast.Print(fset, f)
@@ -103,26 +93,12 @@ func RewriteImportComments(f *ast.File, fset *token.FileSet, prefix string, remo
 		if err != nil {
 			log.Fatalf("Error unquoting import value [%v] - %s\n", parts[1], err)
 		}
-
-		if remove {
-			// the prefix is not there = nothing to remove, keep the comment
-			if !strings.HasPrefix(oldimp, prefix) {
-				newcommentgroups = append(newcommentgroups, c)
-				continue
-			}
-		} else {
-			// the prefix is already in the import path, keep the comment
-			if strings.HasPrefix(oldimp, prefix) {
-				newcommentgroups = append(newcommentgroups, c)
-				continue
-			}
-		}
-		newimp := ""
-		if remove {
-			newimp = oldimp[len(prefix):]
-		} else {
-			newimp = prefix + oldimp
+		// if the prefix is not there = nothing to remove, keep the comment
+		if !strings.HasPrefix(oldimp, from) {
+			newcommentgroups = append(newcommentgroups, c)
+			continue
 		}
+		newimp := strings.Replace(oldimp, from, to, 1)
 		changed = true
 		c2 := ast.Comment{Slash: c.Pos(), Text: `// import ` + strconv.Quote(newimp)}
 		cg := ast.CommentGroup{List: []*ast.Comment{&c2}}

astmod/rewrite_test.go

@@ -11,90 +11,72 @@ import (
 	"testing"
 )
 
-const prefix = `go.corp.example.com/x`
-
-var input map[string]string
-var helloworld, ext1, int1 string
-
-// init reads testdata files and puts them in the input map
-func init() {
-	input = make(map[string]string)
-	base := "testdata/"
-	var err error
-	var b []byte
-	b, err = ioutil.ReadFile(base + "helloworld.go")
-	if err != nil {
-		panic(err)
-	}
-	input["helloworld"] = string(b)
-
-	b, err = ioutil.ReadFile(base + "int1.go")
-	if err != nil {
-		panic(err)
-	}
-	input["int1"] = string(b)
-
-	b, err = ioutil.ReadFile(base + "ext1.go")
-	if err != nil {
-		panic(err)
-	}
-	input["ext1"] = string(b)
-}
+const (
+	from = `go.corp.company.com`
+	to   = `go.newcompany.com`
+)
 
 // tests table
 var tests = []struct {
-	in      string
-	remove  bool
-	wanted  string
+	file    string
+	from    string
+	to      string
+	want    string
 	changed bool
 	label   string
 }{
-	{in: "ext1",
-		remove:  false,
-		wanted:  "int1",
+	{
+		file:    "testdata/ext1.go",
+		from:    "go.corp.example.com",
+		to:      "go.newcompany.com",
+		want:    "testdata/int1.go",
 		changed: true,
 		label:   "ext1",
 	},
-	{in: "int1",
-		remove:  true,
-		wanted:  "ext1",
-		changed: true,
-		label:   "int1",
-	},
-	// try to call rewrite on the file that has already been rewritten - expect a no-op
-	{in: "int1",
-		remove:  false,
-		wanted:  "int1",
+	{
+		file:    "testdata/int1.go",
+		from:    "go.corp.example.com",
+		to:      "go.newcompany.com",
+		want:    "testdata/int1.go",
 		changed: false,
-		label:   "int1-noop",
+		label:   "int1",
 	},
-	{in: "helloworld",
-		remove:  false,
-		wanted:  "helloworld",
+	{
+		file:    "testdata/helloworld.go",
+		from:    "go.corp.example.com",
+		to:      "go.newcompany.com",
+		want:    "testdata/helloworld.go",
 		changed: false,
 		label:   "unmodified",
 	},
 }
 
 func TestRewrite(t *testing.T) {
-	for _, test := range tests {
-		buf, err := Rewrite(test.label, input[test.in], prefix, test.remove)
+	for _, tt := range tests {
+		in, err := ioutil.ReadFile(tt.file)
+		if err != nil {
+			t.Fatal(err)
+		}
+		wantBytes, err := ioutil.ReadFile(tt.want)
+		if err != nil {
+			t.Fatal(err)
+		}
+		want := string(wantBytes)
+		buf, err := Rewrite(tt.label, in, from, to)
 		if err != nil {
 			t.Error(err)
+			continue
 		}
-		// nil buf means no changes - expect test.changed be false
 		if buf == nil {
-			if test.changed == true {
-				t.Errorf("Error in %s - buf is nil but test.changed is true", test.label)
+			if tt.changed == false {
+				continue // OK - no changes were made, as expected
 			}
+			t.Errorf("test '%s': changes expected but none made (buf=nil)", tt.label)
 			continue
 		}
-		if buf != nil && test.changed == false {
-			t.Errorf("Error in %s - buf is non-nil but test.changed is false", test.label)
-		}
-		if buf.String() != input[test.wanted] {
-			t.Errorf("Error in %s: Input:\n%s\n, Got:\n%s\nWanted:\n%s\nremove: %t\n",
-				test.label, input[test.in], buf.String(), input[test.wanted], test.remove)
+		got := buf.String()
+		if got != want {
+			t.Errorf("%s case: got\n%s\nwant contents of %s:\n%s", tt.label, got, tt.want, want)
 		}
 	}
 }

astmod/testdata/ext1.go

@@ -1,9 +1,9 @@
-// package doc comment for github.com/foo/bar
-package bar // import "github.com/foo/bar"
+// package doc comment for go.corp.company.com/abc/xyz
+package bar // import "go.corp.company.com/abc/xyz"
 
-import _ "github.com/abc/xyz"
+import _ "go.corp.company.com/abc/xyz"
 
 func Bar() {
 	// unrelated comment in func
-	println("Hello, World! (from github.com/foo/bar)")
+	println("Hello, World! (from go.corp.company.com/abc/xyz)")
 }

astmod/testdata/helloworld.go

@@ -1,6 +1,10 @@
 package main
 
-import "fmt"
+import (
+	"fmt"
+
+	_ "go.newcompany.com/abc/xyz"
+)
 
 func main() {
 	fmt.Println("Hello, world!")

astmod/testdata/int1.go

@@ -1,9 +1,9 @@
-// package doc comment for github.com/foo/bar
-package bar // import "go.corp.example.com/x/github.com/foo/bar"
+// package doc comment for go.corp.company.com/abc/xyz
+package bar // import "go.newcompany.com/abc/xyz"
 
-import _ "go.corp.example.com/x/github.com/abc/xyz"
+import _ "go.newcompany.com/abc/xyz"
 
 func Bar() {
 	// unrelated comment in func
-	println("Hello, World! (from github.com/foo/bar)")
+	println("Hello, World! (from go.corp.company.com/abc/xyz)")
 }

main.go

@@ -2,10 +2,10 @@
 // Use of this source code is governed by a MIT
 // license that can be found in the LICENSE file.
 //
-// Author: Dmitry Savintsev <dsavints@yahoo-inc.com>
+// Author: Dmitry Savintsev <dsavints@oath.com>
 
 // prewrite tool to rewrite import paths and package import comments for vendoring
-// by adding or removing a given path prefix. The files are rewritten
+// by modifying the given path prefix. The files are rewritten
 // in-place with no backup (expectation is that version control is used), the output is gofmt'ed.
 package main
 
@@ -22,24 +22,27 @@ import (
 )
 
 func usage() {
-	fmt.Fprintf(os.Stderr, "usage: prewrite [flags] [path ...]\n")
+	fmt.Fprintf(os.Stderr, "usage: prewrite [flags] [path]\n")
 	flag.PrintDefaults()
 	os.Exit(2)
 }
 
 func main() {
-	prefix := flag.String("p", "", "package path prefix to prepend to imports (or to remove from imports with -r option)")
-	remove := flag.Bool("r", false, "remove the prefix from import paths")
 	verbose := flag.Bool("v", false, "verbose")
+	from := flag.String("from", "", "package path prefix to replace with the one given in the 'to' parameter")
+	to := flag.String("to", "", "package path prefix to replace the one given in the 'from' parameter")
 	flag.Usage = usage
 	flag.Parse()
-	if *prefix == "" {
+	if *from == "" || *to == "" {
 		usage()
 		os.Exit(1)
 	}
-	// add trailing slash if not already there
-	if (*prefix)[len(*prefix)-1] != '/' {
-		*prefix += "/"
+	// add trailing slashes if not already there
+	if (*from)[len(*from)-1] != '/' {
+		*from += "/"
+	}
+	if (*to)[len(*to)-1] != '/' {
+		*to += "/"
 	}
 	var root string
 	var err error
@@ -51,7 +54,7 @@ func main() {
 	} else {
 		root = flag.Arg(0)
 	}
-	processor := makeVisitor(*prefix, *remove, *verbose)
+	processor := makeVisitor(*from, *to, *verbose)
 	_, err = os.Stat(root)
 	if err != nil && os.IsNotExist(err) {
 		log.Fatalf("Error - the traversal root %s does not exist, please double-check\n", root)
@@ -64,7 +67,7 @@ func main() {
 }
 
 // makeVisitor returns a rewriting function with parameters bound with a closure
-func makeVisitor(prefix string, remove bool, verbose bool) filepath.WalkFunc {
+func makeVisitor(from, to string, verbose bool) filepath.WalkFunc {
 	return func(path string, f os.FileInfo, err error) error {
 		if f.IsDir() || !strings.HasSuffix(f.Name(), ".go") {
 			return nil
@@ -77,7 +80,7 @@ func makeVisitor(prefix string, remove bool, verbose bool) filepath.WalkFunc {
 		if err != nil {
 			log.Fatalf("Fatal error reading file %s\n", path)
 		}
-		buf, err := astmod.Rewrite(path, src, prefix, remove)
+		buf, err := astmod.Rewrite(path, src, from, to)
 		if err != nil {
 			log.Fatalf("Fatal error rewriting AST, file %s - error: %s\n", path, err)
 		}