From 7e9b8d48e50f0c0af5d82f7f310378dfcad9bde2 Mon Sep 17 00:00:00 2001
From: RussellLuo <luopeng.he@gmail.com>
Date: Thu, 5 May 2022 16:45:25 +0800
Subject: [PATCH] pkgtool: Fix PkgPathFromDir and improve PkgNameFromDir

---
 gen/gen.go                  |  5 ++--
 gen/grpc/grpc/generator.go  |  6 ++---
 gen/grpc/proto/generator.go |  6 ++---
 pkg/pkgtool/pkgtool.go      | 53 ++++++++++++++++++++-----------------
 4 files changed, 37 insertions(+), 33 deletions(-)

diff --git a/gen/gen.go b/gen/gen.go
index adcd523..7308425 100644
--- a/gen/gen.go
+++ b/gen/gen.go
@@ -249,8 +249,7 @@ func (g *Generator) generateGRPC(data *ifacetool.Data) (files []*generator.File,
 	if err = ensureDir(pbOutDir); err != nil {
 		return files, err
 	}
-	pbPkgPath := pkgtool.PkgPathFromDir(pbOutDir)
-	f, err := g.proto.Generate(pbPkgPath, data, service)
+	f, err := g.proto.Generate(pbOutDir, data, service)
 	if err != nil {
 		return files, err
 	}
@@ -273,7 +272,7 @@ func (g *Generator) generateGRPC(data *ifacetool.Data) (files []*generator.File,
 
 	// Generate the glue code for adapting the gRPC definition to Go kit.
 	pkgInfo := g.getPkgInfo(outDir)
-	f, err = g.grpc.Generate(pkgInfo, pbPkgPath, data, service)
+	f, err = g.grpc.Generate(pkgInfo, pbOutDir, data, service)
 	if err != nil {
 		return files, err
 	}
diff --git a/gen/grpc/grpc/generator.go b/gen/grpc/grpc/generator.go
index d588d40..16c7712 100644
--- a/gen/grpc/grpc/generator.go
+++ b/gen/grpc/grpc/generator.go
@@ -101,7 +101,7 @@ func New(opts *Options) *Generator {
 	return &Generator{opts: opts}
 }
 
-func (g *Generator) Generate(pkgInfo *generator.PkgInfo, pbPkgPath string, ifaceData *ifacetool.Data, service *parser.Service) (*generator.File, error) {
+func (g *Generator) Generate(pkgInfo *generator.PkgInfo, pbOutDir string, ifaceData *ifacetool.Data, service *parser.Service) (*generator.File, error) {
 	data := struct {
 		PBPkgPath   string
 		PBPkgPrefix string
@@ -109,8 +109,8 @@ func (g *Generator) Generate(pkgInfo *generator.PkgInfo, pbPkgPath string, iface
 		PkgInfo     *generator.PkgInfo
 		Service     *parser.Service
 	}{
-		PBPkgPath:   pbPkgPath,
-		PBPkgPrefix: pkgtool.PkgNameFromDir(pbPkgPath) + ".",
+		PBPkgPath:   pkgtool.PkgPathFromDir(pbOutDir),
+		PBPkgPrefix: pkgtool.PkgNameFromDir(pbOutDir) + ".",
 		Data:        ifaceData,
 		PkgInfo:     pkgInfo,
 		Service:     service,
diff --git a/gen/grpc/proto/generator.go b/gen/grpc/proto/generator.go
index 9c139fd..37c621c 100644
--- a/gen/grpc/proto/generator.go
+++ b/gen/grpc/proto/generator.go
@@ -76,15 +76,15 @@ func New(opts *Options) *Generator {
 	return &Generator{opts: opts}
 }
 
-func (g *Generator) Generate(pkgPath string, ifaceData *ifacetool.Data, service *parser.Service) (*generator.File, error) {
+func (g *Generator) Generate(outDir string, ifaceData *ifacetool.Data, service *parser.Service) (*generator.File, error) {
 	data := struct {
 		PkgPath  string
 		PkgName  string
 		Service  *parser.Service
 		Messages map[string]*parser.Type
 	}{
-		PkgPath:  pkgPath,
-		PkgName:  pkgtool.PkgNameFromDir(pkgPath),
+		PkgPath:  pkgtool.PkgPathFromDir(outDir),
+		PkgName:  pkgtool.PkgNameFromDir(outDir),
 		Service:  service,
 		Messages: getMessages(service),
 	}
diff --git a/pkg/pkgtool/pkgtool.go b/pkg/pkgtool/pkgtool.go
index 55af5ef..908c511 100644
--- a/pkg/pkgtool/pkgtool.go
+++ b/pkg/pkgtool/pkgtool.go
@@ -4,11 +4,9 @@ import (
 	"errors"
 	"fmt"
 	"go/ast"
-	"go/build"
 	"go/parser"
 	"go/token"
 	"path/filepath"
-	"strings"
 
 	"github.com/RussellLuo/kun/pkg/ifacetool"
 	"github.com/RussellLuo/kun/pkg/ifacetool/moq"
@@ -128,39 +126,46 @@ func getInterfaceDoc(ts *ast.TypeSpec, gd *ast.GenDecl) *ast.CommentGroup {
 	return nil
 }
 
-func PkgPathFromDir(dir string) (pkgPath string) {
+func PkgPathFromDir(dir string) string {
 	abs, err := filepath.Abs(dir)
-	if err == nil {
-		pkgPath = stripGopath(abs)
+	if err != nil {
+		panic(err)
 	}
-	return
-}
 
-func PkgNameFromDir(dir string) string {
-	pkg, _ := pkgInfoFromPath(dir, packages.NeedName)
-	if pkg != nil && pkg.Name != "" {
-		return pkg.Name
+	pkg, err := pkgInfoFromPath(abs, packages.NeedModule)
+	if err != nil {
+		panic(err)
 	}
 
-	// Default to the directory name.
-	abs, err := filepath.Abs(dir)
+	if pkg == nil || pkg.Module == nil {
+		return ""
+	}
+
+	// Remove the module root directory.
+	rel, err := filepath.Rel(pkg.Module.Dir, abs)
 	if err != nil {
 		panic(err)
 	}
-	return filepath.Base(abs)
+	// Add the module path prefix.
+	return filepath.Join(pkg.Module.Path, rel)
 }
 
-// stripGopath takes the directory to a package and removes the
-// $GOPATH/src path to get the canonical package name.
-func stripGopath(p string) string {
-	for _, srcDir := range build.Default.SrcDirs() {
-		rel, err := filepath.Rel(srcDir, p)
-		if err != nil || strings.HasPrefix(rel, "..") {
-			continue
-		}
-		return filepath.ToSlash(rel)
+func PkgNameFromDir(dir string) string {
+	abs, err := filepath.Abs(dir)
+	if err != nil {
+		panic(err)
+	}
+
+	pkg, err := pkgInfoFromPath(abs, packages.NeedName)
+	if err != nil {
+		panic(err)
+	}
+	if pkg != nil && pkg.Name != "" {
+		return pkg.Name
 	}
-	return p
+
+	// Default to the directory name.
+	return filepath.Base(abs)
 }
 
 func pkgInfoFromPath(srcDir string, mode packages.LoadMode) (*packages.Package, error) {