From 2d508d3a5187c74010cacb5c59f0d408f9f6beaa Mon Sep 17 00:00:00 2001
From: adrianiacobghiula <2491756+adrianiacobghiula@users.noreply.github.com>
Date: Mon, 6 May 2024 19:11:50 +0200
Subject: [PATCH] allow custom template

---
 cmd/avrogen/main.go      |  4 ++--
 gen/gen.go               | 19 ++++++++-----------
 gen/output_template.tmpl |  3 +--
 go.mod                   |  2 ++
 go.sum                   |  4 ++++
 5 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/cmd/avrogen/main.go b/cmd/avrogen/main.go
index bb7b05e..90c51dc 100644
--- a/cmd/avrogen/main.go
+++ b/cmd/avrogen/main.go
@@ -5,7 +5,6 @@ import (
 	"errors"
 	"flag"
 	"fmt"
-	"go/format"
 	"io"
 	"os"
 	"path/filepath"
@@ -13,6 +12,7 @@ import (
 
 	"github.com/hamba/avro/v2"
 	"github.com/hamba/avro/v2/gen"
+	"golang.org/x/tools/imports"
 )
 
 type config struct {
@@ -94,7 +94,7 @@ func realMain(args []string, stdout, stderr io.Writer) int {
 		_, _ = fmt.Fprintf(stderr, "Error: could not generate code: %v\n", err)
 		return 3
 	}
-	formatted, err := format.Source(buf.Bytes())
+	formatted, err := imports.Process("", buf.Bytes(), nil)
 	if err != nil {
 		_, _ = fmt.Fprintf(stderr, "Error: generated code \n%s\n could not be formatted: %v\n", buf.String(), err)
 		return 3
diff --git a/gen/gen.go b/gen/gen.go
index 70c101d..2f5e1e7 100644
--- a/gen/gen.go
+++ b/gen/gen.go
@@ -6,13 +6,14 @@ import (
 	_ "embed"
 	"errors"
 	"fmt"
-	"go/format"
 	"io"
+	"maps"
 	"strings"
 	"text/template"
 
 	"github.com/ettle/strcase"
 	"github.com/hamba/avro/v2"
+	"golang.org/x/tools/imports"
 )
 
 // Config configures the code generation.
@@ -82,7 +83,7 @@ func StructFromSchema(schema avro.Schema, w io.Writer, cfg Config) error {
 		return err
 	}
 
-	formatted, err := format.Source(buf.Bytes())
+	formatted, err := imports.Process("", buf.Bytes(), nil)
 	if err != nil {
 		return fmt.Errorf("generated code \n%s\n could not be formatted: %w", buf.String(), err)
 	}
@@ -149,10 +150,13 @@ type Generator struct {
 
 // NewGenerator returns a generator.
 func NewGenerator(pkg string, tags map[string]TagStyle, opts ...OptsFunc) *Generator {
+	clonedTags := maps.Clone(tags)
+	delete(clonedTags, "avro")
+
 	g := &Generator{
 		template: outputTemplate,
 		pkg:      pkg,
-		tags:     tags,
+		tags:     clonedTags,
 	}
 
 	for _, opt := range opts {
@@ -295,7 +299,6 @@ func (g *Generator) resolveLogicalSchema(logicalType avro.LogicalType) string {
 }
 
 func (g *Generator) newField(name, typ, avroFieldDoc, avroFieldName string) field {
-	delete(g.tags, "avro")
 	return field{
 		Name:          name,
 		Type:          typ,
@@ -337,12 +340,6 @@ func (g *Generator) Write(w io.Writer) error {
 		return err
 	}
 
-	imports := g.imports
-	if len(g.thirdPartyImports) > 0 {
-		imports = append(imports, "")
-		imports = append(imports, g.thirdPartyImports...)
-	}
-
 	data := struct {
 		WithEncoders      bool
 		PackageName       string
@@ -352,7 +349,7 @@ func (g *Generator) Write(w io.Writer) error {
 	}{
 		WithEncoders: g.encoders,
 		PackageName:  g.pkg,
-		Imports:      imports,
+		Imports:      append(g.imports, g.thirdPartyImports...),
 		Typedefs:     g.typedefs,
 	}
 	return parsed.Execute(w, data)
diff --git a/gen/output_template.tmpl b/gen/output_template.tmpl
index b81cb2d..94cbcf6 100644
--- a/gen/output_template.tmpl
+++ b/gen/output_template.tmpl
@@ -6,8 +6,7 @@ package {{ .PackageName }}
 {{ if len .Imports }}
 	import (
 	{{- range .Imports }}
-		{{ if eq . "" }}
-		{{ else }}"{{ . }}"{{end}}
+		"{{ . }}"
 	{{- end }}
 	)
 {{ end }}
diff --git a/go.mod b/go.mod
index 4c1a3bd..4586f11 100644
--- a/go.mod
+++ b/go.mod
@@ -16,5 +16,7 @@ require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	golang.org/x/mod v0.17.0 // indirect
+	golang.org/x/tools v0.20.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
diff --git a/go.sum b/go.sum
index 99fc2b0..3db61f1 100644
--- a/go.sum
+++ b/go.sum
@@ -23,6 +23,10 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+golang.org/x/mod v0.17.0 h1:zY54UmvipHiNd+pm+m0x9KhZ9hl1/7QNMyxXbc6ICqA=
+golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
+golang.org/x/tools v0.20.0 h1:hz/CVckiOxybQvFw6h7b/q80NTr9IUQb4s1IIzW7KNY=
+golang.org/x/tools v0.20.0/go.mod h1:WvitBU7JJf6A4jOdg4S1tviW9bhUxkgeCui/0JHctQg=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=