diff --git a/cmd/hz/app/app.go b/cmd/hz/app/app.go index 116df93cc..739db8823 100644 --- 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 index 4e6d75ed5..48e381297 100644 --- 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 index d6f775cef..a503965a6 100644 --- 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, "/") {