From 7934811009abfc75dfe5cc36c1dfc11bc6d9d43c Mon Sep 17 00:00:00 2001
From: twacqwq <xiaot.qaq@gmail.com>
Date: Fri, 30 Jun 2023 01:21:21 +0800
Subject: [PATCH] feat(cmd): set custom module name for command new

---
 cmd/kratos/internal/base/mod.go        | 19 +++++++++++++++++++
 cmd/kratos/internal/project/new.go     | 14 +++++++++++---
 cmd/kratos/internal/project/project.go | 15 ++++++++++-----
 3 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/cmd/kratos/internal/base/mod.go b/cmd/kratos/internal/base/mod.go
index 2349aca9e..378e1377f 100644
--- a/cmd/kratos/internal/base/mod.go
+++ b/cmd/kratos/internal/base/mod.go
@@ -60,3 +60,22 @@ func KratosMod() string {
 	// $GOPATH/src/github.com/go-kratos/kratos
 	return filepath.Join(gopath, "src", "github.com", "go-kratos", "kratos")
 }
+
+// ModuleName returns custom module name
+func ModuleName(moduleFile, moduleName string) error {
+	modBytes, err := os.ReadFile(moduleFile)
+	if err != nil {
+		return err
+	}
+	goMod, err := modfile.Parse(moduleFile, modBytes, nil)
+	if err != nil {
+		return err
+	}
+	goMod.Module.Syntax.Token[1] = moduleName
+	modBytes, err = goMod.Format()
+	if err != nil {
+		return err
+	}
+
+	return os.WriteFile(moduleFile, modBytes, 0644)
+}
diff --git a/cmd/kratos/internal/project/new.go b/cmd/kratos/internal/project/new.go
index 038b22049..ba4b1aced 100644
--- a/cmd/kratos/internal/project/new.go
+++ b/cmd/kratos/internal/project/new.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"path"
 	"path/filepath"
 
 	"github.com/AlecAivazis/survey/v2"
@@ -14,8 +15,9 @@ import (
 
 // Project is a project template.
 type Project struct {
-	Name string
-	Path string
+	Name       string
+	Path       string
+	ModuleName string
 }
 
 // New new a project from remote repo.
@@ -37,7 +39,7 @@ func (p *Project) New(ctx context.Context, dir string, layout string, branch str
 		}
 		os.RemoveAll(to)
 	}
-	fmt.Printf("🚀 Creating service %s, layout repo is %s, please wait a moment.\n\n", p.Name, layout)
+	fmt.Printf("🚀 Creating service %s, layout repo is %s, module name is %s, please wait a moment.\n\n", p.Name, layout, p.ModuleName)
 	repo := base.NewRepo(layout, branch)
 	if err := repo.CopyTo(ctx, to, p.Name, []string{".git", ".github"}); err != nil {
 		return err
@@ -49,6 +51,12 @@ func (p *Project) New(ctx context.Context, dir string, layout string, branch str
 	if e != nil {
 		return e
 	}
+	if p.ModuleName != "" && p.ModuleName != p.Name {
+		e = base.ModuleName(path.Join(to, "go.mod"), p.ModuleName)
+		if e != nil {
+			return e
+		}
+	}
 	base.Tree(to, dir)
 
 	fmt.Printf("\n🍺 Project creation succeeded %s\n", color.GreenString(p.Name))
diff --git a/cmd/kratos/internal/project/project.go b/cmd/kratos/internal/project/project.go
index ef3bcad80..3a7625cc8 100644
--- a/cmd/kratos/internal/project/project.go
+++ b/cmd/kratos/internal/project/project.go
@@ -24,10 +24,11 @@ var CmdNew = &cobra.Command{
 }
 
 var (
-	repoURL string
-	branch  string
-	timeout string
-	nomod   bool
+	repoURL    string
+	branch     string
+	timeout    string
+	nomod      bool
+	moduleName string
 )
 
 func init() {
@@ -39,6 +40,7 @@ func init() {
 	CmdNew.Flags().StringVarP(&branch, "branch", "b", branch, "repo branch")
 	CmdNew.Flags().StringVarP(&timeout, "timeout", "t", timeout, "time out")
 	CmdNew.Flags().BoolVarP(&nomod, "nomod", "", nomod, "retain go mod")
+	CmdNew.Flags().StringVarP(&moduleName, "module-name", "m", moduleName, "module name")
 }
 
 func run(_ *cobra.Command, args []string) {
@@ -64,9 +66,12 @@ func run(_ *cobra.Command, args []string) {
 		}
 	} else {
 		name = args[0]
+		if moduleName == "" {
+			moduleName = name
+		}
 	}
 	projectName, workingDir := processProjectParams(name, wd)
-	p := &Project{Name: projectName}
+	p := &Project{Name: projectName, ModuleName: moduleName}
 	done := make(chan error, 1)
 	go func() {
 		if !nomod {