minor refactoring and added comments

2 files changed, 46 insertions, 43 deletions

diff --git a/ast.go b/ast.go
index beb6a1a..0a79e41 100644
--- a/ast.go
+++ b/ast.go
@@ -8,7 +8,6 @@ import (
 	"go/token"
 	"io/ioutil"
 	"os"
-	"path"
 	"strings"
 
 	"github.com/mgutz/ansi"
@@ -25,56 +24,53 @@ func ProcessFileAST(filePath string, from string, to string) {
 	//Reset the color
 	reset := ansi.ColorCode("reset")
 
-	// If the file is a go file scan it
-	if path.Ext(filePath) == ".go" {
-		// New FileSet to parse the go file to
-		fSet := token.NewFileSet()
+	// New FileSet to parse the go file to
+	fSet := token.NewFileSet()
 
-		// Parse the file
-		file, err := parser.ParseFile(fSet, filePath, nil, 0)
-		if err != nil {
-			fmt.Println(err)
-		}
+	// Parse the file
+	file, err := parser.ParseFile(fSet, filePath, nil, 0)
+	if err != nil {
+		fmt.Println(err)
+	}
 
-		// Get the list of imports from the ast
-		imports := astutil.Imports(fSet, file)
+	// Get the list of imports from the ast
+	imports := astutil.Imports(fSet, file)
 
-		// Keep track of number of changes
-		numChanges := 0
+	// 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 := mImport.Path.Value
+	// Iterate through the imports array
+	for _, mPackage := range imports {
+		for _, mImport := range mPackage {
+			// Since astutil returns the path string with quotes, remove those
+			importString := 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++
+			// 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)
+				// 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 " + importString + " from file " + reset + white + filePath + reset + red + " to " + reset + white + replacePackage + reset)
+				fmt.Println(red + "Updating import " + importString + " from file " + reset + white + filePath + 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)
-				}
+				// Remove the old import and replace it with the replacement
+				astutil.DeleteImport(fSet, file, importString)
+				astutil.AddImport(fSet, file, replacePackage)
 			}
 		}
+	}
 
-		// If the number of changes are more than 0, write file
-		if numChanges > 0 {
-			// Print the new AST tree to a new output buffer
-			var outputBuffer bytes.Buffer
-			printer.Fprint(&outputBuffer, fSet, file)
+	// If the number of changes are more than 0, write file
+	if numChanges > 0 {
+		// Print the new AST tree to a new output buffer
+		var outputBuffer bytes.Buffer
+		printer.Fprint(&outputBuffer, fSet, file)
 
-			ioutil.WriteFile(filePath, outputBuffer.Bytes(), os.ModePerm)
-			fmt.Printf(blackOnWhite+"File "+filePath+" saved after %d changes."+reset+"\n", numChanges)
-		} else {
-			fmt.Println(greenUnderline + "No changes needed on file " + filePath + reset)
-		}
+		ioutil.WriteFile(filePath, outputBuffer.Bytes(), os.ModePerm)
+		fmt.Printf(blackOnWhite+"File "+filePath+" saved after %d changes."+reset+"\n", numChanges)
+	} else {
+		fmt.Println(greenUnderline + "No changes needed on file " + filePath + reset)
 	}
 }
diff --git a/main.go b/main.go
index e7c51df..3d5d856 100644
--- a/main.go
+++ b/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"os"
+	"path"
 	"path/filepath"
 
 	"github.com/codegangsta/cli"
@@ -36,7 +37,7 @@ func main() {
 		if file != "" {
 			ProcessFileNative(file, from, to)
 		} else {
-			RunApp(dir, from, to, c)
+			ScanDir(dir, from, to, c)
 		}
 
 	}
@@ -44,11 +45,17 @@ func main() {
 	app.Run(os.Args)
 }
 
-func RunApp(dir string, from string, to string, c *cli.Context) {
+// ScanDir scans a directory for go files and
+func ScanDir(dir string, from string, to string, c *cli.Context) {
 
+	// If from and to are not empty scan all files
 	if from != "" && to != "" {
+		// Scan directory for files
 		filepath.Walk(dir, func(filePath string, info os.FileInfo, err error) error {
-			ProcessFileNative(filePath, from, to)
+			// Only process go files
+			if path.Ext(filePath) == ".go" {
+				ProcessFileNative(filePath, from, to)
+			}
 			return nil
 		})