From deacbef89190fb9708154b6c281e3f946ae9078a Mon Sep 17 00:00:00 2001
From: tbxark <tbxark@outlook.com>
Date: Sat, 9 Nov 2024 19:04:11 +0800
Subject: [PATCH] =?UTF-8?q?chore:=20=E6=95=B4=E7=90=86protoc=E6=8F=92?=
 =?UTF-8?q?=E4=BB=B6=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 contrib/protoc-gen-route/route.go | 41 ++++++++++++++++++++-----------
 contrib/protoc-gen-sphere/http.go | 29 +++++++++++++---------
 2 files changed, 45 insertions(+), 25 deletions(-)

diff --git a/contrib/protoc-gen-route/route.go b/contrib/protoc-gen-route/route.go
index 027fd35..1f54144 100644
--- 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,16 +44,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
-		}
-		routeTemplate = string(raw)
-	}
-	generateFileContent(gen, file, g, conf)
-	return g
 }
 
 func generateFileContent(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile, conf *Config) {
@@ -53,6 +51,14 @@ func generateFileContent(gen *protogen.Plugin, file *protogen.File, g *protogen.
 		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
index 5cb945b..5b4cf19 100644
--- 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)
+	}
+}