Handle package aliases correctly

acceptance/fixtures/fakes/named_package_interface.go

@@ -0,0 +1,31 @@
+package fakes
+
+import (
+	"sync"
+
+	"github.com/cloudfoundry/bosh-utils/logger"
+	clogger "github.com/hashicorp/consul/logger"
+)
+
+type NamedPackageInterface struct {
+	NamedPackageMethodCall struct {
+		sync.Mutex
+		CallCount int
+		Receives  struct {
+			Config clogger.Config
+			Level  logger.LogLevel
+		}
+		Stub func(clogger.Config, logger.LogLevel)
+	}
+}
+
+func (f *NamedPackageInterface) NamedPackageMethod(param1 clogger.Config, param2 logger.LogLevel) {
+	f.NamedPackageMethodCall.Lock()
+	defer f.NamedPackageMethodCall.Unlock()
+	f.NamedPackageMethodCall.CallCount++
+	f.NamedPackageMethodCall.Receives.Config = param1
+	f.NamedPackageMethodCall.Receives.Level = param2
+	if f.NamedPackageMethodCall.Stub != nil {
+		f.NamedPackageMethodCall.Stub(param1, param2)
+	}
+}

acceptance/fixtures/garbage/interface.go

@@ -1 +1,3 @@
+package garbage
+
 garbage

acceptance/fixtures/interfaces.go

@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"io"
 
+	"github.com/cloudfoundry/bosh-utils/logger"
+	clogger "github.com/hashicorp/consul/logger"
 	"github.com/pivotal-cf/jhanda"
 )
 
@@ -31,3 +33,7 @@ type DuplicateArgumentInterface interface {
 type FunctionInterface interface {
 	FuncMethod(func(string) error) func(int) bool
 }
+
+type NamedPackageInterface interface {
+	NamedPackageMethod(config clogger.Config, level logger.LogLevel)
+}

acceptance/generate_test.go

@@ -64,6 +64,7 @@ var _ = Describe("faux", func() {
 		Entry("stdlib", "", "io", "Reader", "io_reader.go"),
 		Entry("variadic", "./fixtures/interfaces.go", "", "VariadicInterface", "variadic_interface.go"),
 		Entry("functions", "./fixtures/interfaces.go", "", "FunctionInterface", "function_interface.go"),
+		Entry("named package", "./fixtures/interfaces.go", "", "NamedPackageInterface", "named_package_interface.go"),
 	)
 
 	Context("when the source file is provided via an environment variable", func() {

go.mod

@@ -1,10 +1,13 @@
 module github.com/ryanmoran/faux
 
 require (
+	github.com/cloudfoundry/bosh-utils v0.0.0-20191005100153-94a8490a77f0
+	github.com/hashicorp/consul v1.6.1
 	github.com/onsi/ginkgo v1.7.0
 	github.com/onsi/gomega v1.4.3
 	github.com/pivotal-cf/jhanda v0.0.0-20181025233525-e6aa09a032df
-	golang.org/x/tools v0.0.0-20190228203856-589c23e65e65
+	golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135
+	google.golang.org/appengine v1.4.0 // indirect
 )
 
 go 1.13

main.go

@@ -98,7 +98,7 @@ Flags:
 		stderr.Fatalf("could not determine output absolute path: %s", err)
 	}
 
-	result, err := imports.Process(outputPath, buffer.Bytes(), &imports.Options{FormatOnly: false})
+	result, err := imports.Process(outputPath, buffer.Bytes(), &imports.Options{FormatOnly: true})
 	if err != nil {
 		stderr.Fatalf("could not process imports: %s\n\n%s", err, buffer.String())
 	}

parsing/argument.go

@@ -8,14 +8,28 @@ type Argument struct {
 	Name     string
 	Type     types.Type
 	Variadic bool
-	Package  string
+	Package  *types.Package
 }
 
-func NewArgument(v *types.Var, variadic bool) Argument {
-	var pkg string
-	if t, ok := v.Type().(*types.Named); ok {
+func NewArgument(pkgMap map[string]string, v *types.Var, variadic bool) Argument {
+	var pkg *types.Package
+	switch t := v.Type().(type) {
+	case *types.Named:
 		if t.Obj().Pkg() != nil {
-			pkg = t.Obj().Pkg().Path()
+			pkg = t.Obj().Pkg()
+		}
+	case *types.Pointer:
+		if e, ok := t.Elem().(*types.Named); ok {
+			if e.Obj().Pkg() != nil {
+				pkg = e.Obj().Pkg()
+			}
+		}
+	}
+
+	if pkg != nil {
+		name := pkgMap[pkg.Path()]
+		if name != "" {
+			pkg.SetName(name)
 		}
 	}
 

parsing/interface.go

@@ -10,7 +10,7 @@ type Interface struct {
 	Signatures []Signature
 }
 
-func NewInterface(n *types.Named) (Interface, error) {
+func NewInterface(pkgMap map[string]string, n *types.Named) (Interface, error) {
 	var signatures []Signature
 
 	underlying, ok := n.Underlying().(*types.Interface)
@@ -19,7 +19,7 @@ func NewInterface(n *types.Named) (Interface, error) {
 	}
 
 	for i := 0; i < underlying.NumMethods(); i++ {
-		signatures = append(signatures, NewSignature(underlying.Method(i)))
+		signatures = append(signatures, NewSignature(pkgMap, underlying.Method(i)))
 	}
 
 	return Interface{

parsing/interface_test.go

@@ -25,7 +25,7 @@ var _ = Describe("Interface", func() {
 	})
 
 	It("parses an interface from a named type", func() {
-		iface, err := parsing.NewInterface(namedType)
+		iface, err := parsing.NewInterface(nil, namedType)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(iface).To(Equal(parsing.Interface{
 			Name: "SomeType",
@@ -44,7 +44,7 @@ var _ = Describe("Interface", func() {
 		})
 
 		It("includes those methods in the parsed interface", func() {
-			iface, err := parsing.NewInterface(namedType)
+			iface, err := parsing.NewInterface(nil, namedType)
 			Expect(err).NotTo(HaveOccurred())
 			Expect(iface).To(Equal(parsing.Interface{
 				Name: "SomeType",
@@ -64,7 +64,7 @@ var _ = Describe("Interface", func() {
 		})
 
 		It("returns an error", func() {
-			_, err := parsing.NewInterface(namedType)
+			_, err := parsing.NewInterface(nil, namedType)
 			Expect(err).To(MatchError("failed to load underlying type: int is not an interface"))
 		})
 	})

parsing/parse.go

@@ -2,7 +2,11 @@ package parsing
 
 import (
 	"fmt"
+	"go/parser"
+	"go/token"
 	"go/types"
+	"path/filepath"
+	"strconv"
 
 	"golang.org/x/tools/go/packages"
 )
@@ -18,6 +22,35 @@ func Parse(path, name string) (Interface, error) {
 	}
 	pkg := pkgs[0].Types
 
+	if len(pkgs[0].GoFiles) == 0 {
+		return Interface{}, fmt.Errorf("failed to load package with any files: %q", path)
+	}
+
+	dir := filepath.Dir(pkgs[0].GoFiles[0])
+	astPkgs, err := parser.ParseDir(token.NewFileSet(), dir, nil, parser.ImportsOnly)
+	if err != nil {
+		return Interface{}, fmt.Errorf("failed to parse directory: %s", dir)
+	}
+
+	pkgMap := map[string]string{}
+	for _, pkg := range astPkgs {
+		for _, file := range pkg.Files {
+			for _, i := range file.Imports {
+				var name string
+				if i.Name != nil {
+					name = i.Name.Name
+				}
+
+				path, err := strconv.Unquote(i.Path.Value)
+				if err != nil {
+					return Interface{}, err
+				}
+
+				pkgMap[path] = name
+			}
+		}
+	}
+
 	object := pkg.Scope().Lookup(name)
 	if object == nil {
 		return Interface{}, fmt.Errorf("failed to find named type: %s.%s", path, name)
@@ -28,5 +61,5 @@ func Parse(path, name string) (Interface, error) {
 		return Interface{}, fmt.Errorf("failed to load named type: %s.%s", path, name)
 	}
 
-	return NewInterface(namedType)
+	return NewInterface(pkgMap, namedType)
 }

parsing/parse_test.go

@@ -97,9 +97,16 @@ var _ = Describe("Parse", func() {
 		})
 
 		Context("when the name matches no object in scope", func() {
+			It("returns an error", func() {
+				_, err := parsing.Parse("io", "Banana")
+				Expect(err).To(MatchError("failed to find named type: io.Banana"))
+			})
+		})
+
+		Context("when the package has no files", func() {
 			It("returns an error", func() {
 				_, err := parsing.Parse("some-package", "SomeType")
-				Expect(err).To(MatchError("failed to find named type: some-package.SomeType"))
+				Expect(err).To(MatchError("failed to load package with any files: \"some-package\""))
 			})
 		})
 	})

parsing/signature.go

@@ -8,7 +8,7 @@ type Signature struct {
 	Results []Argument
 }
 
-func NewSignature(f *types.Func) Signature {
+func NewSignature(pkgMap map[string]string, f *types.Func) Signature {
 	signature := f.Type().(*types.Signature)
 
 	var params []Argument
@@ -18,12 +18,12 @@ func NewSignature(f *types.Func) Signature {
 			variadic = true
 		}
 
-		params = append(params, NewArgument(signature.Params().At(i), variadic))
+		params = append(params, NewArgument(pkgMap, signature.Params().At(i), variadic))
 	}
 
 	var results []Argument
 	for i := 0; i < signature.Results().Len(); i++ {
-		results = append(results, NewArgument(signature.Results().At(i), false))
+		results = append(results, NewArgument(pkgMap, signature.Results().At(i), false))
 	}
 
 	return Signature{

parsing/signature_test.go

@@ -23,7 +23,7 @@ var _ = Describe("Signature", func() {
 	It("parses a signature from a func", func() {
 		method := types.NewFunc(0, pkg, "SomeMethod", sig)
 
-		signature := parsing.NewSignature(method)
+		signature := parsing.NewSignature(nil, method)
 		Expect(signature).To(Equal(parsing.Signature{
 			Name: "SomeMethod",
 		}))
@@ -47,7 +47,7 @@ var _ = Describe("Signature", func() {
 		It("parses those params", func() {
 			method := types.NewFunc(0, pkg, "SomeMethod", sig)
 
-			signature := parsing.NewSignature(method)
+			signature := parsing.NewSignature(nil, method)
 			Expect(signature).To(Equal(parsing.Signature{
 				Name: "SomeMethod",
 				Params: []parsing.Argument{
@@ -73,7 +73,7 @@ var _ = Describe("Signature", func() {
 			It("marks the last param as variadic", func() {
 				method := types.NewFunc(0, pkg, "SomeMethod", sig)
 
-				signature := parsing.NewSignature(method)
+				signature := parsing.NewSignature(nil, method)
 				Expect(signature).To(Equal(parsing.Signature{
 					Name: "SomeMethod",
 					Params: []parsing.Argument{

rendering/context.go

@@ -6,10 +6,16 @@ import (
 	"github.com/ryanmoran/faux/parsing"
 )
 
-type Context struct{}
+type Context struct {
+	imports Imports
+}
 
 func NewContext() *Context {
-	return &Context{}
+	return &Context{
+		imports: Imports{
+			Package{Path: "sync"},
+		},
+	}
 }
 
 func (c *Context) Build(iface parsing.Interface) File {
@@ -20,7 +26,7 @@ func (c *Context) Build(iface parsing.Interface) File {
 		funcs = append(funcs, c.BuildFunc(fake, signature))
 	}
 
-	return NewFile("fakes", []NamedType{fake}, funcs)
+	return NewFile("fakes", c.imports, []NamedType{fake}, funcs)
 }
 
 func (c *Context) BuildFakeType(iface parsing.Interface) NamedType {
@@ -63,6 +69,10 @@ func (c *Context) BuildCallCount() Field {
 func (c *Context) BuildReceives(args []parsing.Argument) Field {
 	var fields []Field
 	for i, arg := range args {
+		if arg.Package != nil {
+			c.imports.Add(arg.Package)
+		}
+
 		name := arg.Name
 		if name == "" {
 			name = FieldTypeName(args, i)
@@ -78,6 +88,10 @@ func (c *Context) BuildReceives(args []parsing.Argument) Field {
 func (c *Context) BuildReturns(args []parsing.Argument) Field {
 	var fields []Field
 	for i, arg := range args {
+		if arg.Package != nil {
+			c.imports.Add(arg.Package)
+		}
+
 		name := arg.Name
 		if name == "" {
 			name = FieldTypeName(args, i)

rendering/file.go

@@ -1,16 +1,23 @@
 package rendering
 
-import "go/ast"
+import (
+	"go/ast"
+	"go/token"
+
+	"golang.org/x/tools/go/ast/astutil"
+)
 
 type File struct {
 	Package string
+	Imports Imports
 	Types   []NamedType
 	Funcs   []Func
 }
 
-func NewFile(pkg string, types []NamedType, funcs []Func) File {
+func NewFile(pkg string, imports Imports, types []NamedType, funcs []Func) File {
 	return File{
 		Package: pkg,
+		Imports: imports,
 		Types:   types,
 		Funcs:   funcs,
 	}
@@ -26,8 +33,14 @@ func (f File) AST() *ast.File {
 		decls = append(decls, fn.Decl())
 	}
 
-	return &ast.File{
+	file := &ast.File{
 		Name:  ast.NewIdent(f.Package),
 		Decls: decls,
 	}
+
+	for _, pkg := range f.Imports {
+		astutil.AddNamedImport(token.NewFileSet(), file, pkg.Name, pkg.Path)
+	}
+
+	return file
 }