fix pb type (#57)

pubgo · Dec 5, 2024 · af18b82 · af18b82
1 parent dc12759
commit af18b82
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git a/cmds/protoc-gen-lava/internal/gen.go b/cmds/protoc-gen-lava/internal/gen.go
@@ -73,16 +73,11 @@ func GenerateFile(gen *protogen.Plugin, file *protogen.File) (g *protogen.Genera
 		keyPrefix := strings.ReplaceAll(srvInfo.GoName, "InnerService", "")
 		keyPrefix = strings.ReplaceAll(keyPrefix, "Inner", "") + "Service"
 
-		var pbPkg = ""
-		if file.GoImportPath != meta.mth.Input.GoIdent.GoImportPath {
-			pbPkg = string(meta.mth.Input.GoIdent.GoImportPath)
-		}
-
 		genFile.Var().Op("_").Op("=").Qual(rpcMetaPkg, "Register").Call(
 			jen.Op("&").Qual(rpcMetaPkg, "RpcMeta").Values(
 				jen.Dict{
-					jen.Id("Input"):  jen.New(jen.Qual(pbPkg, meta.mth.Input.GoIdent.GoName)),
-					jen.Id("Output"): jen.New(jen.Qual(pbPkg, meta.mth.Output.GoIdent.GoName)),
+					jen.Id("Input"):  jen.New(getPkg(file, meta.mth.Input.GoIdent)),
+					jen.Id("Output"): jen.New(getPkg(file, meta.mth.Output.GoIdent)),
 					jen.Id("Name"):   jen.Lit(meta.meta.Name),
 					jen.Id("Method"): jen.Lit(fmt.Sprintf("%s/%s", meta.srv.Desc.FullName(), meta.mth.GoName)),
 					jen.Id("Tags"): jen.Map(jen.String()).String().Values(
@@ -99,3 +94,12 @@ func GenerateFile(gen *protogen.Plugin, file *protogen.File) (g *protogen.Genera
 	g.P(genFile.GoString())
 	return g
 }
+
+func getPkg(file *protogen.File, goIdent protogen.GoIdent) *jen.Statement {
+	var pkgName = ""
+	if file.GoImportPath != goIdent.GoImportPath {
+		pkgName = string(goIdent.GoImportPath)
+	}
+
+	return jen.Qual(pkgName, goIdent.GoName)
+}
diff --git a/cmds/protoc-gen-lava/internal/version.go b/cmds/protoc-gen-lava/internal/version.go
@@ -2,11 +2,11 @@ package internal
 
 import (
 	"fmt"
-	
+
 	"google.golang.org/protobuf/compiler/protogen"
 )
 
-const Version = "v0.0.2"
+const Version = "v0.0.3"
 
 func protocVersion(gen *protogen.Plugin) string {
 	v := gen.Request.GetCompilerVersion()