chore: 整理protoc插件代码

TBXark · Nov 9, 2024 · deacbef · deacbef
1 parent ee0d6f4
commit deacbef
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 25 deletions.
diff --git a/contrib/protoc-gen-route/route.go b/contrib/protoc-gen-route/route.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"github.com/tbxark/options-proto/go/tbxark/options"
 	"google.golang.org/protobuf/compiler/protogen"
 	"google.golang.org/protobuf/proto"
@@ -26,6 +27,13 @@ func generateFile(gen *protogen.Plugin, file *protogen.File, conf *Config) *prot
 	}
 	filename := file.GeneratedFilenamePrefix + conf.genFileSuffix
 	g := gen.NewGeneratedFile(filename, file.GoImportPath)
+	genBaseFileHeader(g, file)
+	replaceTemplateIfNeed(conf)
+	generateFileContent(gen, file, g, conf)
+	return g
+}
+
+func genBaseFileHeader(g *protogen.GeneratedFile, file *protogen.File) {
 	g.P("// Code generated by protoc-gen-route. DO NOT EDIT.")
 	g.P("// versions:")
 	if file.Proto.GetOptions().GetDeprecated() {
@@ -36,23 +44,21 @@ func generateFile(gen *protogen.Plugin, file *protogen.File, conf *Config) *prot
 	g.P()
 	g.P("package ", file.GoPackageName)
 	g.P()
-	if conf.templateFile != "" {
-		raw, err := os.ReadFile(conf.templateFile)
-		if err != nil {
-			gen.Error(err)
-			return nil
-		}
-		routeTemplate = string(raw)
-	}
-	generateFileContent(gen, file, g, conf)
-	return g
 }
 
 func generateFileContent(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile, conf *Config) {
 	if len(file.Services) == 0 {
 		return
 	}
 	g.P("var _ = new(", contextPackage.Ident("Context"), ")")
+	genGoImports(g, conf)
+	g.P()
+	for _, service := range file.Services {
+		genService(gen, file, g, service, conf)
+	}
+}
+
+func genGoImports(g *protogen.GeneratedFile, conf *Config) {
 	idents := []*GoIdent{
 		conf.requestType,
 		conf.responseType,
@@ -71,10 +77,6 @@ func generateFileContent(gen *protogen.Plugin, file *protogen.File, g *protogen.
 			imported[string(i.pkg)] = struct{}{}
 		}
 	}
-	g.P()
-	for _, service := range file.Services {
-		genService(gen, file, g, service, conf)
-	}
 }
 
 func genService(_ *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile, service *protogen.Service, conf *Config) {
@@ -144,6 +146,17 @@ func extractOptionsRule(method *protogen.Method, key string) *options.KeyValuePa
 	return nil
 }
 
+func replaceTemplateIfNeed(conf *Config) {
+	if conf.templateFile != "" {
+		raw, err := os.ReadFile(conf.templateFile)
+		if err != nil {
+			_, _ = fmt.Fprintf(os.Stderr, "read template file error: %v\n", err)
+			os.Exit(2)
+		}
+		routeTemplate = string(raw)
+	}
+}
+
 func pascalCase(s string) string {
 	words := strings.FieldsFunc(s, func(r rune) bool {
 		return r == ' ' || r == '_' || r == '-'

diff --git a/contrib/protoc-gen-sphere/http.go b/contrib/protoc-gen-sphere/http.go
@@ -40,6 +40,13 @@ func generateFile(gen *protogen.Plugin, file *protogen.File, conf *Config) *prot
 	}
 	filename := file.GeneratedFilenamePrefix + "_sphere.pb.go"
 	g := gen.NewGeneratedFile(filename, file.GoImportPath)
+	genBaseFileHeader(gen, file, g)
+	replaceTemplateIfNeed(conf)
+	generateFileContent(gen, file, g, conf)
+	return g
+}
+
+func genBaseFileHeader(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile) {
 	g.P("// Code generated by protoc-gen-sphere. DO NOT EDIT.")
 	g.P("// versions:")
 	g.P("// - protoc             ", protocVersion(gen))
@@ -51,17 +58,6 @@ func generateFile(gen *protogen.Plugin, file *protogen.File, conf *Config) *prot
 	g.P()
 	g.P("package ", file.GoPackageName)
 	g.P()
-
-	if conf.templateFile != "" {
-		raw, err := os.ReadFile(conf.templateFile)
-		if err != nil {
-			gen.Error(err)
-			return nil
-		}
-		httpTemplate = string(raw)
-	}
-	generateFileContent(gen, file, g, conf)
-	return g
 }
 
 // generateFileContent generates the sphere errors definitions, excluding the package statement.
@@ -512,3 +508,14 @@ func newGenConf(g *protogen.GeneratedFile, conf *Config) *genConfig {
 	}
 	return genConf
 }
+
+func replaceTemplateIfNeed(conf *Config) {
+	if conf.templateFile != "" {
+		raw, err := os.ReadFile(conf.templateFile)
+		if err != nil {
+			_, _ = fmt.Fprintf(os.Stderr, "read template file error: %v\n", err)
+			os.Exit(2)
+		}
+		httpTemplate = string(raw)
+	}
+}