Keep comments in safe mode.

2 files changed, 25 insertions, 43 deletions

diff --git a/README.md b/README.md
index cbd0d93..da599fc 100644
--- a/README.md
+++ b/README.md
@@ -57,7 +57,7 @@ You can also run `gomove --help` for help.
 
 Safe Mode
 -------------
-By default this tool uses something called native mode where it scans through the go files, finds import statements and replaces the package names. However if you want to use a more safe way of changing the import statements, you can use Safe Mode which parses the ast tree of the go file and replaces imports from there. The biggest downside to using this is that it will remove all the comments from go files and also prettify your files. To use it in safe mode set the `--safe-mode` or `-s` flag to true:
+By default this tool uses something called native mode where it scans through the go files, finds import statements and replaces the package names. However if you want to use a more safe way of changing the import statements, you can use Safe Mode which parses the ast tree of the go file and replaces imports from there. The biggest downside to using this is that it will [prettify](https://godoc.org/github.com/golang/go/src/cmd/gofmt) your files. To use it in safe mode set the `--safe-mode` or `-s` flag to true:
     
     gomove --safe-mode true github.com/bla/bla github.com/foo/bar
 
diff --git a/ast.go b/ast.go
index 5259379..41ab3df 100644
--- a/ast.go
+++ b/ast.go
@@ -8,7 +8,6 @@ import (
 	"go/token"
 	"io/ioutil"
 	"os"
-	"strings"
 
 	"github.com/mgutz/ansi"
 	"golang.org/x/tools/go/ast/astutil"
@@ -31,61 +30,44 @@ func ProcessFileAST(filePath string, from string, to string) {
 	fSet := token.NewFileSet()
 
 	// Parse the file
-	file, err := parser.ParseFile(fSet, filePath, nil, 0)
+	file, err := parser.ParseFile(fSet, filePath, nil, parser.ParseComments)
 	if err != nil {
 		fmt.Println(err)
+		return
 	}
 
-	// Get the list of imports from the ast
-	imports := astutil.Imports(fSet, file)
-
 	// Keep track of number of changes
-	numChanges := 0
-
-	// Iterate through the imports array
-	for _, mPackage := range imports {
-		for _, mImport := range mPackage {
-			// Since astutil returns the path string with quotes, remove those
-			importString := strings.TrimSuffix(strings.TrimPrefix(mImport.Path.Value, "\""), "\"")
-
-			// If the path matches the oldpath, replace it with the new one
-			if strings.Contains(importString, from) {
-				//If it needs to be replaced, increase numChanges so we can write the file later
-				numChanges++
-
-				// Join the path of the import package with the remainder from the old one after removing the old import package
-				replacePackage := strings.Replace(importString, from, to, -1)
-
-				fmt.Println(red +
-					"Updating import " +
-					reset + white +
-					importString +
-					reset + red +
-					" to " +
-					reset + white +
-					replacePackage +
-					reset)
-
-				// Remove the old import and replace it with the replacement
-				astutil.DeleteImport(fSet, file, importString)
-				astutil.AddImport(fSet, file, replacePackage)
-			}
-		}
+	changed := false
+	if changed = astutil.RewriteImport(fSet, file, from, to); changed {
+		fmt.Println(red +
+			"Updating import " +
+			reset + white +
+			from +
+			reset + red +
+			" to " +
+			reset + white +
+			to +
+			reset)
 	}
 
 	// If the number of changes are more than 0, write file
-	if numChanges > 0 {
-		// Print the new AST tree to a new output buffer
+	if changed {
+		// Print the new AST tree to a new output buffer. These Config settings intended to match gofmt.
+		printerMode := printer.TabIndent | printer.UseSpaces
+		printConfig := &printer.Config{Mode: printerMode, Tabwidth: 8}
+
 		var outputBuffer bytes.Buffer
-		printer.Fprint(&outputBuffer, fSet, file)
+		err := printConfig.Fprint(&outputBuffer, fSet, file)
+		if err != nil {
+			fmt.Println(err)
+			return
+		}
 
 		ioutil.WriteFile(filePath, outputBuffer.Bytes(), os.ModePerm)
 		fmt.Println(yellow+
 			"File",
 			filePath,
-			"saved after",
-			numChanges,
-			"changes",
+			"saved",
 			reset, "\n\n")
 	} else {
 		fmt.Println(yellow+