From 7d704b6b0dfd098a64e1daa2d31f4b3f75eb9a2b Mon Sep 17 00:00:00 2001
From: Max Chechel <hexdigest@gmail.com>
Date: Sun, 25 Aug 2024 16:21:57 -0400
Subject: [PATCH] Fixes #87

---
 cmd_generate.go        |  5 ++++-
 generator/generator.go | 17 +++++++++++------
 go.mod                 |  1 +
 go.sum                 |  2 ++
 4 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/cmd_generate.go b/cmd_generate.go
index d809f5b..4fa137c 100644
--- a/cmd_generate.go
+++ b/cmd_generate.go
@@ -307,5 +307,8 @@ package {{.Package.Name}}
 {{if (not .Options.HeaderVars.DisableGoGenerate)}}
 //{{"go:generate"}} gowrap gen -p {{.SourcePackage.PkgPath}} -i {{.Options.InterfaceName}} -t {{.Options.HeaderVars.Template}} -o {{.Options.HeaderVars.OutputFileName}}{{.Options.HeaderVars.VarsArgs}} -l "{{.Options.LocalPrefix}}"
 {{end}}
-
+import(
+{{range $import := .Options.Imports}}{{$import}}
+	{{end}}
+)
 `
diff --git a/generator/generator.go b/generator/generator.go
index 0856b30..cfcd463 100644
--- a/generator/generator.go
+++ b/generator/generator.go
@@ -13,6 +13,8 @@ import (
 	"text/template"
 
 	"github.com/pkg/errors"
+	"golang.org/x/text/cases"
+	"golang.org/x/text/language"
 	"golang.org/x/tools/go/packages"
 	"golang.org/x/tools/imports"
 
@@ -199,16 +201,19 @@ func NewGenerator(options Options) (*Generator, error) {
 		return nil, errors.Wrap(err, "failed to parse source package")
 	}
 
-	interfaceType := srcPackage.Name + "." + options.InterfaceName
+	var interfaceType string
 	if srcPackage.PkgPath == dstPackage.PkgPath {
 		interfaceType = options.InterfaceName
 		srcPackageAST.Name = ""
 	} else {
 		if options.SourcePackageAlias != "" {
 			srcPackageAST.Name = options.SourcePackageAlias
+		} else {
+			srcPackageAST.Name = "_source" + cases.Title(language.Und, cases.NoLower).String(srcPackageAST.Name)
 		}
 
-		options.Imports = append(options.Imports, `"`+srcPackage.PkgPath+`"`)
+		interfaceType = srcPackageAST.Name + "." + options.InterfaceName
+		options.Imports = append(options.Imports, srcPackageAST.Name+` "`+srcPackage.PkgPath+`"`)
 	}
 
 	output, err := findTarget(processInput{
@@ -251,12 +256,12 @@ func NewGenerator(options Options) (*Generator, error) {
 
 func makeImports(imports []*ast.ImportSpec) []string {
 	result := make([]string, len(imports))
-	for _, i := range imports {
+	for i, im := range imports {
 		var name string
-		if i.Name != nil {
-			name = i.Name.Name
+		if im.Name != nil {
+			name = im.Name.Name
 		}
-		result = append(result, name+" "+i.Path.Value)
+		result[i] = name + " " + im.Path.Value
 	}
 
 	return result
diff --git a/go.mod b/go.mod
index abc472b..49af979 100644
--- a/go.mod
+++ b/go.mod
@@ -7,6 +7,7 @@ require (
 	github.com/gojuno/minimock/v3 v3.0.10
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.8.4
+	golang.org/x/text v0.3.7
 	golang.org/x/tools v0.1.12
 )
 
diff --git a/go.sum b/go.sum
index 91c520f..1fc31f7 100644
--- a/go.sum
+++ b/go.sum
@@ -138,6 +138,8 @@ golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
+golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=