From 8aa8c525b82c3762c4ee9a308e34e7e5099aa8f2 Mon Sep 17 00:00:00 2001
From: Daniel Podolsky <1049874+onokonem@users.noreply.github.com>
Date: Mon, 30 Sep 2024 12:40:37 +0300
Subject: [PATCH] minor fixes

---
 README.md                                     |  2 +-
 .../internal/flags/root.go                    |  8 +----
 internal/generator/generator.go               | 34 ++++++++++++++-----
 3 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/README.md b/README.md
index 062bc0a..1aea5bf 100644
--- a/README.md
+++ b/README.md
@@ -18,7 +18,7 @@ Please rise the issue in case of any difference found in these plugins output.
 ### Install
 
 ```
-go install github.com/Djarvur/protoc-gen-python-grpc/cmd/protoc-gen-python-grpc@v0.1.3
+go install github.com/Djarvur/protoc-gen-python-grpc/cmd/protoc-gen-python-grpc@v0.1.4
 ```
 
 Make sure your [protoc](https://grpc.io/docs/protoc-installation/)/[buf](https://buf.build/docs/installation) compiler can see the `protoc-gen-python-grpc` in the path.
diff --git a/cmd/protoc-gen-python-grpc/internal/flags/root.go b/cmd/protoc-gen-python-grpc/internal/flags/root.go
index c84b117..1c205b6 100644
--- a/cmd/protoc-gen-python-grpc/internal/flags/root.go
+++ b/cmd/protoc-gen-python-grpc/internal/flags/root.go
@@ -35,13 +35,7 @@ func Root() *cobra.Command {
 }
 
 func runRoot(suffix, templateSource string) {
-	err := protokit.RunPlugin(
-		&generator.Generator{
-			Suffix:   suffix,
-			Template: templateSource,
-		},
-	)
-	if err != nil {
+	if err := protokit.RunPlugin(must(generator.New(suffix, templateSource))); err != nil {
 		panic(err)
 	}
 }
diff --git a/internal/generator/generator.go b/internal/generator/generator.go
index 8df153d..081c628 100644
--- a/internal/generator/generator.go
+++ b/internal/generator/generator.go
@@ -37,16 +37,30 @@ type Method struct {
 // SupportedFeatures describes a flag setting for supported features.
 const SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)
 
-// Generator describes a protoc code generate plugin.
-// It's an implementation of Generator from github.com/pseudomuto/protokit.
-type Generator struct {
+var _ protokit.Plugin = (*generator)(nil)
+
+// generator describes a protoc code generate plugin.
+// It's an implementation of generator from github.com/pseudomuto/protokit.
+type generator struct {
 	Suffix   string
-	Template string
+	Template *template.Template
+}
+
+func New(suffix, tmplSrc string) (*generator, error) {
+	tmpl, err := buildTemplate(tmplSrc)
+	if err != nil {
+		return nil, err
+	}
+
+	return &generator{
+		Suffix:   suffix,
+		Template: tmpl,
+	}, nil
 }
 
 // Generate compiles the documentation and generates the CodeGeneratorResponse to send back to protoc. It does this
 // by rendering a template based on the options parsed from the CodeGeneratorRequest.
-func (p *Generator) Generate(r *pluginpb.CodeGeneratorRequest) (*pluginpb.CodeGeneratorResponse, error) {
+func (p *generator) Generate(r *pluginpb.CodeGeneratorRequest) (*pluginpb.CodeGeneratorResponse, error) {
 	resp := new(pluginpb.CodeGeneratorResponse)
 
 	for _, fds := range protokit.ParseCodeGenRequest(r) {
@@ -75,7 +89,7 @@ func (p *Generator) Generate(r *pluginpb.CodeGeneratorRequest) (*pluginpb.CodeGe
 	return resp, nil
 }
 
-func executeTemplate(tmplSrc string, data interface{}) (string, error) {
+func buildTemplate(tmplSrc string) (*template.Template, error) {
 	tmplFuncs := template.FuncMap{
 		"trimSuffix": strings.TrimSuffix,
 		"baseName":   strings.BaseName,
@@ -86,12 +100,16 @@ func executeTemplate(tmplSrc string, data interface{}) (string, error) {
 
 	tmpl, err := template.New("").Funcs(tmplFuncs).Parse(tmplSrc)
 	if err != nil {
-		return "", fmt.Errorf("parsing template: %w", err)
+		return nil, fmt.Errorf("parsing template: %w", err)
 	}
 
+	return tmpl, nil
+}
+
+func executeTemplate(tmpl *template.Template, data interface{}) (string, error) {
 	buf := new(bytes.Buffer)
 
-	if err = tmpl.Execute(buf, data); err != nil {
+	if err := tmpl.Execute(buf, data); err != nil {
 		return "", fmt.Errorf("executing template: %w", err)
 	}