fixed merged imports, was not including sql drivers for example

.travis.yml

@@ -12,7 +12,6 @@ script:
 
   # pull all external dependencies
   - go get github.com/cucumber/gherkin-go
-  - go get golang.org/x/tools/imports
   - go get github.com/shiena/ansicolor
 
   # run standard go tests

Makefile

@@ -11,5 +11,4 @@ test:
 # updates dependencies
 deps:
 	go get -u github.com/cucumber/gherkin-go
-	go get -u golang.org/x/tools/imports
 	go get -u github.com/shiena/ansicolor

builder.go

@@ -8,10 +8,9 @@ import (
 	"go/token"
 	"os"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"text/template"
-
-	"golang.org/x/tools/imports"
 )
 
 type builder struct {
@@ -20,6 +19,24 @@ type builder struct {
 	Contexts []string
 	Internal bool
 	tpl      *template.Template
+
+	imports []*ast.ImportSpec
+}
+
+func (b *builder) hasImport(a *ast.ImportSpec) bool {
+	for _, b := range b.imports {
+		var aname, bname string
+		if a.Name != nil {
+			aname = a.Name.Name
+		}
+		if b.Name != nil {
+			bname = b.Name.Name
+		}
+		if bname == aname && a.Path.Value == b.Path.Value {
+			return true
+		}
+	}
+	return false
 }
 
 func newBuilder(buildPath string) (*builder, error) {
@@ -40,7 +57,7 @@ func main() {
 }`)),
 	}
 
-	return b, filepath.Walk(buildPath, func(path string, file os.FileInfo, err error) error {
+	err := filepath.Walk(buildPath, func(path string, file os.FileInfo, err error) error {
 		if file.IsDir() && file.Name() != "." {
 			return filepath.SkipDir
 		}
@@ -51,6 +68,8 @@ func main() {
 		}
 		return err
 	})
+
+	return b, err
 }
 
 func (b *builder) parseFile(path string) error {
@@ -66,20 +85,24 @@ func (b *builder) parseFile(path string) error {
 	b.registerSteps(f)
 	b.deleteImports(f)
 	b.files[path] = f
+
 	return nil
 }
 
 func (b *builder) deleteImports(f *ast.File) {
 	var decls []ast.Decl
 	for _, d := range f.Decls {
-		fun, ok := d.(*ast.GenDecl)
-		if !ok {
-			decls = append(decls, d)
+		gen, ok := d.(*ast.GenDecl)
+		if ok && gen.Tok == token.IMPORT {
+			for _, spec := range gen.Specs {
+				impspec := spec.(*ast.ImportSpec)
+				if !b.hasImport(impspec) {
+					b.imports = append(b.imports, impspec)
+				}
+			}
 			continue
 		}
-		if fun.Tok != token.IMPORT {
-			decls = append(decls, fun)
-		}
+		decls = append(decls, d)
 	}
 	f.Decls = decls
 }
@@ -132,13 +155,38 @@ func (b *builder) merge() (*ast.File, error) {
 	if err != nil {
 		return nil, err
 	}
+	// b.imports(f)
 	b.deleteImports(f)
 	b.files["main.go"] = f
 
 	pkg, _ := ast.NewPackage(b.fset, b.files, nil, nil)
 	pkg.Name = "main"
 
-	return ast.MergePackageFiles(pkg, ast.FilterImportDuplicates), nil
+	ret, err := ast.MergePackageFiles(pkg, 0), nil
+	if err != nil {
+		return ret, err
+	}
+
+	// @TODO: we reread the file, probably something goes wrong with position
+	buf.Reset()
+	if err = format.Node(&buf, b.fset, ret); err != nil {
+		return nil, err
+	}
+
+	ret, err = parser.ParseFile(b.fset, "", buf.Bytes(), 0)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, spec := range b.imports {
+		var name string
+		if spec.Name != nil {
+			name = spec.Name.Name
+		}
+		ipath, _ := strconv.Unquote(spec.Path.Value)
+		addImport(b.fset, ret, name, ipath)
+	}
+	return ret, nil
 }
 
 // Build creates a runnable Godog executable file
@@ -168,5 +216,147 @@ func Build() ([]byte, error) {
 		return nil, err
 	}
 
-	return imports.Process("", buf.Bytes(), nil)
+	return buf.Bytes(), nil
+}
+
+// taken from https://github.com/golang/tools/blob/master/go/ast/astutil/imports.go#L17
+func addImport(fset *token.FileSet, f *ast.File, name, ipath string) {
+	newImport := &ast.ImportSpec{
+		Path: &ast.BasicLit{
+			Kind:  token.STRING,
+			Value: strconv.Quote(ipath),
+		},
+	}
+	if name != "" {
+		newImport.Name = &ast.Ident{Name: name}
+	}
+
+	// Find an import decl to add to.
+	// The goal is to find an existing import
+	// whose import path has the longest shared
+	// prefix with ipath.
+	var (
+		bestMatch  = -1         // length of longest shared prefix
+		lastImport = -1         // index in f.Decls of the file's final import decl
+		impDecl    *ast.GenDecl // import decl containing the best match
+		impIndex   = -1         // spec index in impDecl containing the best match
+	)
+	for i, decl := range f.Decls {
+		gen, ok := decl.(*ast.GenDecl)
+		if ok && gen.Tok == token.IMPORT {
+			lastImport = i
+			// Do not add to import "C", to avoid disrupting the
+			// association with its doc comment, breaking cgo.
+			if declImports(gen, "C") {
+				continue
+			}
+
+			// Match an empty import decl if that's all that is available.
+			if len(gen.Specs) == 0 && bestMatch == -1 {
+				impDecl = gen
+			}
+
+			// Compute longest shared prefix with imports in this group.
+			for j, spec := range gen.Specs {
+				impspec := spec.(*ast.ImportSpec)
+				n := matchLen(importPath(impspec), ipath)
+				if n > bestMatch {
+					bestMatch = n
+					impDecl = gen
+					impIndex = j
+				}
+			}
+		}
+	}
+
+	// If no import decl found, add one after the last import.
+	if impDecl == nil {
+		impDecl = &ast.GenDecl{
+			Tok: token.IMPORT,
+		}
+		if lastImport >= 0 {
+			impDecl.TokPos = f.Decls[lastImport].End()
+		} else {
+			// There are no existing imports.
+			// Our new import goes after the package declaration and after
+			// the comment, if any, that starts on the same line as the
+			// package declaration.
+			impDecl.TokPos = f.Package
+
+			file := fset.File(f.Package)
+			pkgLine := file.Line(f.Package)
+			for _, c := range f.Comments {
+				if file.Line(c.Pos()) > pkgLine {
+					break
+				}
+				impDecl.TokPos = c.End()
+			}
+		}
+		f.Decls = append(f.Decls, nil)
+		copy(f.Decls[lastImport+2:], f.Decls[lastImport+1:])
+		f.Decls[lastImport+1] = impDecl
+	}
+
+	// Insert new import at insertAt.
+	insertAt := 0
+	if impIndex >= 0 {
+		// insert after the found import
+		insertAt = impIndex + 1
+	}
+	impDecl.Specs = append(impDecl.Specs, nil)
+	copy(impDecl.Specs[insertAt+1:], impDecl.Specs[insertAt:])
+	impDecl.Specs[insertAt] = newImport
+	pos := impDecl.Pos()
+	if insertAt > 0 {
+		// Assign same position as the previous import,
+		// so that the sorter sees it as being in the same block.
+		pos = impDecl.Specs[insertAt-1].Pos()
+	}
+	if newImport.Name != nil {
+		newImport.Name.NamePos = pos
+	}
+	newImport.Path.ValuePos = pos
+	newImport.EndPos = pos
+
+	// Clean up parens. impDecl contains at least one spec.
+	if len(impDecl.Specs) == 1 {
+		// Remove unneeded parens.
+		impDecl.Lparen = token.NoPos
+	} else if !impDecl.Lparen.IsValid() {
+		// impDecl needs parens added.
+		impDecl.Lparen = impDecl.Specs[0].Pos()
+	}
+
+	f.Imports = append(f.Imports, newImport)
+}
+
+func declImports(gen *ast.GenDecl, path string) bool {
+	if gen.Tok != token.IMPORT {
+		return false
+	}
+	for _, spec := range gen.Specs {
+		impspec := spec.(*ast.ImportSpec)
+		if importPath(impspec) == path {
+			return true
+		}
+	}
+	return false
+}
+
+func matchLen(x, y string) int {
+	n := 0
+	for i := 0; i < len(x) && i < len(y) && x[i] == y[i]; i++ {
+		if x[i] == '/' {
+			n++
+		}
+	}
+	return n
+}
+
+func importPath(s *ast.ImportSpec) string {
+	t, err := strconv.Unquote(s.Path.Value)
+	if err == nil {
+		return t
+	}
+	return ""
 }

cmd/godog/main.go

@@ -1,10 +1,8 @@
 package main
 
 import (
-	"fmt"
 	"os"
 	"os/exec"
-	"time"
 
 	"github.com/DATA-DOG/godog"
 	"github.com/shiena/ansicolor"
@@ -14,12 +12,11 @@ func buildAndRun() error {
 	// will support Ansi colors for windows
 	stdout := ansicolor.NewAnsiColorWriter(os.Stdout)
 
-	builtFile := fmt.Sprintf("%s/%dgodog.go", os.TempDir(), time.Now().UnixNano())
+	builtFile := "/tmp/bgodog.go"
 	// @TODO: then there is a suite error or panic, it may
 	// be interesting to see the built file. But we
 	// even cannot determine the status of exit error
 	// so leaving it for the future
-	defer os.Remove(builtFile)
 
 	buf, err := godog.Build()
 	if err != nil {
@@ -29,6 +26,7 @@ func buildAndRun() error {
 	if err != nil {
 		return err
 	}
+	defer os.Remove(builtFile)
 	if _, err = w.Write(buf); err != nil {
 		w.Close()
 		return err