feat(hz): trim pb gopackage (#1111)

cloudwego · Jun 6, 2024 · 7437071 · 7437071
1 parent d273c77
commit 7437071
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 8 deletions.
diff --git a/cmd/hz/app/app.go b/cmd/hz/app/app.go
@@ -191,6 +191,7 @@ func Init() *cli.App {
 	customLayoutData := cli.StringFlag{Name: "customize_layout_data_path", Usage: "Specify the path for layout template render data.", Destination: &globalArgs.CustomizeLayoutData}
 	customPackage := cli.StringFlag{Name: "customize_package", Usage: "Specify the path for package template.", Destination: &globalArgs.CustomizePackage}
 	handlerByMethod := cli.BoolFlag{Name: "handler_by_method", Usage: "Generate a separate handler file for each method.", Destination: &globalArgs.HandlerByMethod}
+	trimGoPackage := cli.StringFlag{Name: "trim_gopackage", Aliases: []string{"trim_pkg"}, Usage: "Trim the prefix of go_package for protobuf.", Destination: &globalArgs.TrimGoPackage}
 
 	// app
 	app := cli.NewApp()
@@ -225,6 +226,7 @@ func Init() *cli.App {
 				&thriftOptionsFlag,
 				&protoOptionsFlag,
 				&optPkgFlag,
+				&trimGoPackage,
 				&noRecurseFlag,
 				&forceNewFlag,
 				&enableExtendsFlag,
@@ -261,6 +263,7 @@ func Init() *cli.App {
 				&thriftOptionsFlag,
 				&protoOptionsFlag,
 				&optPkgFlag,
+				&trimGoPackage,
 				&noRecurseFlag,
 				&enableExtendsFlag,
 				&sortRouterFlag,
@@ -291,6 +294,7 @@ func Init() *cli.App {
 				&thriftOptionsFlag,
 				&protoOptionsFlag,
 				&noRecurseFlag,
+				&trimGoPackage,
 
 				&jsonEnumStrFlag,
 				&unsetOmitemptyFlag,
@@ -318,6 +322,7 @@ func Init() *cli.App {
 				&protoOptionsFlag,
 				&noRecurseFlag,
 				&enableExtendsFlag,
+				&trimGoPackage,
 
 				&jsonEnumStrFlag,
 				&queryEnumIntFlag,

diff --git a/cmd/hz/config/argument.go b/cmd/hz/config/argument.go
@@ -41,12 +41,13 @@ type Argument struct {
 	BaseDomain     string // request domain
 	ForceClientDir string // client dir (not use namespace as a subpath)
 
-	IdlType   string   // idl type
-	IdlPaths  []string // master idl path
-	RawOptPkg []string // user-specified package import path
-	OptPkgMap map[string]string
-	Includes  []string
-	PkgPrefix string
+	IdlType       string   // idl type
+	IdlPaths      []string // master idl path
+	RawOptPkg     []string // user-specified package import path
+	OptPkgMap     map[string]string
+	Includes      []string
+	PkgPrefix     string
+	TrimGoPackage string // trim go_package for protobuf, avoid to generate multiple directory
 
 	Gopath      string // $GOPATH
 	Gosrc       string // $GOPATH/src

diff --git a/cmd/hz/protobuf/plugin.go b/cmd/hz/protobuf/plugin.go
@@ -198,7 +198,7 @@ func (plugin *Plugin) Response(resp *pluginpb.CodeGeneratorResponse) error {
 }
 
 func (plugin *Plugin) Handle(req *pluginpb.CodeGeneratorRequest, args *config.Argument) error {
-	plugin.fixGoPackage(req, plugin.PkgMap)
+	plugin.fixGoPackage(req, plugin.PkgMap, args.TrimGoPackage)
 
 	// new plugin
 	opts := protogen.Options{}
@@ -291,12 +291,16 @@ func (plugin *Plugin) Handle(req *pluginpb.CodeGeneratorRequest, args *config.Ar
 }
 
 // fixGoPackage will update go_package to store all the model files in ${model_dir}
-func (plugin *Plugin) fixGoPackage(req *pluginpb.CodeGeneratorRequest, pkgMap map[string]string) {
+func (plugin *Plugin) fixGoPackage(req *pluginpb.CodeGeneratorRequest, pkgMap map[string]string, trimGoPackage string) {
 	gopkg := plugin.Package
 	for _, f := range req.ProtoFile {
 		if strings.HasPrefix(f.GetPackage(), "google.protobuf") {
 			continue
 		}
+		if len(trimGoPackage) != 0 && strings.HasPrefix(f.GetOptions().GetGoPackage(), trimGoPackage) {
+			*f.Options.GoPackage = strings.TrimPrefix(*f.Options.GoPackage, trimGoPackage)
+		}
+
 		opt := getGoPackage(f, pkgMap)
 		if !strings.Contains(opt, gopkg) {
 			if strings.HasPrefix(opt, "/") {