Merge pull request #10 from k-yomo/fix-missing-package-name

Fix import bug happened when the package name is same as th target package

cmd/fixtory/main.go

@@ -53,9 +53,9 @@ func main() {
 		outputPath = filepath.Join(targetDir, "fixtory_gen.go")
 	}
 
+	outputDir, _ := path.Split(outputPath)
 	newWriter := func() (io.Writer, func(), error) {
-		dir, _ := path.Split(outputPath)
-		if err := os.MkdirAll(dir, 0755); err != nil {
+		if err := os.MkdirAll(outputDir, 0755); err != nil {
 			return nil, nil, xerrors.Errorf("create directory: %w", err)
 		}
 		writer, err := os.Create(outputPath)
@@ -64,7 +64,7 @@ func main() {
 		}
 		return writer, func() { _ = writer.Close() }, nil
 	}
-	if err := fixtory.Generate(targetDir, types, *pkgName, newWriter); err != nil {
+	if err := fixtory.Generate(targetDir, filepath.Dir(outputDir), types, *pkgName, newWriter); err != nil {
 		color.Red("%+v", err)
 		os.Exit(1)
 	}

factory_test.go

@@ -1,9 +1,10 @@
 package fixtory
 
 import (
-	"github.com/google/go-cmp/cmp"
 	"reflect"
 	"testing"
+
+	"github.com/google/go-cmp/cmp"
 )
 
 type testStruct struct {
@@ -89,7 +90,7 @@ func TestBuilder_Build(t *testing.T) {
 			want:    &testStruct{Int: 5},
 		},
 		{
-			name: "struct is overwritten by traits, zero, each param",
+			name:    "struct is overwritten by traits, zero, each param",
 			builder: fac.NewBuilder(bluePrint, testStruct{String: "setByTrait1", Int: 10}, testStruct{String: "setByTrait2", Array: []int{1, 2, 3}}).Zero("Map").EachParam(testStruct{Float: 10.9}),
 			want: &testStruct{
 				String: "setByTrait2",
@@ -104,7 +105,7 @@ func TestBuilder_Build(t *testing.T) {
 			},
 		},
 		{
-			name: "empty fields do not overwrite",
+			name:    "empty fields do not overwrite",
 			builder: fac.NewBuilder(bluePrint, testStruct{}).EachParam(testStruct{}),
 			want: &testStruct{
 				String: "setByBlueprint",

generate.go

@@ -3,28 +3,41 @@ package fixtory
 import (
 	"bytes"
 	"fmt"
-	"github.com/k-yomo/fixtory/pkg/astutil"
 	"go/ast"
 	"go/format"
-	"golang.org/x/xerrors"
 	"io"
 	"strings"
 	"text/template"
+
+	"github.com/k-yomo/fixtory/pkg/astutil"
+	"golang.org/x/xerrors"
 )
 
-func Generate(targetDir string, types []string, pkgName string, newWriter func() (writer io.Writer, close func(), err error)) error {
+func Generate(targetDir string, outputDir string, types []string, pkgName string, newWriter func() (writer io.Writer, close func(), err error)) error {
 	targetTypeMap := map[string]bool{}
 	for _, t := range types {
 		targetTypeMap[t] = true
 	}
+	if len(targetTypeMap) == 0 {
+		return nil
+	}
 
 	walkerMap, err := astutil.DirToAstWalker(targetDir)
 	if err != nil {
 		return err
 	}
+	if len(walkerMap) == 0 {
+		return nil
+	}
 	for _, walker := range walkerMap {
-		if len(targetTypeMap) == 0 {
-			break
+		if pkgName == "" {
+			pkgName = walker.Pkg.Name
+		}
+
+		importPkgName := walker.Pkg.Name
+		shouldImportPkg := outputDir != targetDir
+		if shouldImportPkg && pkgName == walker.Pkg.Name {
+			importPkgName = fmt.Sprintf("_%s", importPkgName)
 		}
 		body := new(bytes.Buffer)
 		for _, spec := range walker.AllStructSpecs() {
@@ -43,16 +56,16 @@ func Generate(targetDir string, types []string, pkgName string, newWriter func()
 			}
 			tpl := template.Must(template.New("factory").Funcs(template.FuncMap{"ToLower": strings.ToLower}).Parse(factoryTpl))
 			st := spec.Name.Name
-			if pkgName != "" {
-				st = fmt.Sprintf("%s.%s", walker.Pkg.Name, st)
+			if shouldImportPkg {
+				st = fmt.Sprintf("%s.%s", importPkgName, st)
 			}
 			params := struct {
 				StructName string
-				Struct string
+				Struct     string
 				FieldNames []string
 			}{
 				StructName: spec.Name.Name,
-				Struct: st,
+				Struct:     st,
 				FieldNames: fieldNames,
 			}
 			if err := tpl.Execute(body, params); err != nil {
@@ -64,10 +77,12 @@ func Generate(targetDir string, types []string, pkgName string, newWriter func()
 		}
 
 		var importPackages []string
-		if pkgName == "" {
-			pkgName = walker.Pkg.Name
-		} else {
-			importPackages = append(importPackages, walker.PkgPath)
+		if shouldImportPkg {
+			if pkgName == walker.Pkg.Name {
+				importPackages = append(importPackages, fmt.Sprintf(`%s "%s"`, importPkgName, walker.PkgPath))
+			} else {
+				importPackages = append(importPackages, fmt.Sprintf(`"%s"`, walker.PkgPath))
+			}
 		}
 
 		out := new(bytes.Buffer)

generate_test.go

@@ -9,6 +9,7 @@ import (
 func TestGenerate(t *testing.T) {
 	type args struct {
 		targetDir string
+		outputDir string
 		types     []string
 		pkgName   string
 		newWriter func() (io.Writer, func(), error)
@@ -22,17 +23,18 @@ func TestGenerate(t *testing.T) {
 			name: "generate factory for Author, Article in example directory",
 			args: args{
 				targetDir: "example",
+				outputDir: "example",
 				types:     []string{"Article", "Author"},
 				newWriter: func() (io.Writer, func(), error) {
 					var b bytes.Buffer
-					return &b, func() { }, nil
+					return &b, func() {}, nil
 				},
 			},
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			if err := Generate(tt.args.targetDir, tt.args.types, tt.args.pkgName, tt.args.newWriter); (err != nil) != tt.wantErr {
+			if err := Generate(tt.args.targetDir, tt.args.outputDir, tt.args.types, tt.args.pkgName, tt.args.newWriter); (err != nil) != tt.wantErr {
 				t.Errorf("Generate() error = %v, wantErr %v", err, tt.wantErr)
 			}
 		})