Repository: hamba/avro
Pull request #290 (2023-08-19 15:24:13)

fix: refactor struct generator
This refactors the struct generator.


Patch:
diff --git a/gen/gen.go b/gen/gen.go
index 5ee539a..2ac8019 100644
--- a/gen/gen.go
+++ b/gen/gen.go
@@ -35,8 +35,9 @@ const (
 	UpperCamel TagStyle = "upper-camel"
 )
 
-const outputTemplate = `// Code generated by avro/gen. DO NOT EDIT.
-package {{ .PackageName }}
+const outputTemplate = `package {{ .PackageName }}
+
+// Code generated by avro/gen. DO NOT EDIT.
 
 {{ if len .Imports }}
 import (
@@ -92,7 +93,10 @@ func StructFromSchema(schema avro.Schema, w io.Writer, cfg Config) error {
 		return errors.New("can only generate Go code from Record Schemas")
 	}
 
-	g := NewGenerator(strcase.ToSnake(cfg.PackageName), cfg.Tags, WithFullName(cfg.FullName))
+	opts := []OptsFunc{
+		WithFullName(cfg.FullName),
+	}
+	g := NewGenerator(strcase.ToSnake(cfg.PackageName), cfg.Tags, opts...)
 	g.Parse(rec)
 
 	buf := &bytes.Buffer{}
@@ -145,36 +149,24 @@ func NewGenerator(pkg string, tags map[string]TagStyle, opts ...OptsFunc) *Gener
 	return g
 }
 
+// Reset reset the generator.
+func (g *Generator) Reset() {
+	g.imports = g.imports[:0]
+	g.thirdPartyImports = g.thirdPartyImports[:0]
+	g.typedefs = g.typedefs[:0]
+}
+
 // Parse parses an avro schema into Go types.
 func (g *Generator) Parse(schema avro.Schema) {
 	_ = g.generate(schema)
 }
 
 func (g *Generator) generate(schema avro.Schema) string {
-	switch t := schema.(type) {
-	case *avro.RecordSchema:
-		typeName := g.resolveTypeName(t)
-		fields := make([]field, len(t.Fields()))
-		for i, f := range t.Fields() {
-			fSchema := f.Type()
-			fieldName := strcase.ToGoPascal(f.Name())
-			typ := g.resolveType(fSchema)
-			tag := f.Name()
-			fields[i] = g.newField(fieldName, typ, tag)
-		}
-		g.typedefs = append(g.typedefs, newType(typeName, fields))
-		return typeName
-	default:
-		return g.resolveType(schema)
-	}
-}
-
-func (g *Generator) resolveType(schema avro.Schema) string {
 	switch s := schema.(type) {
 	case *avro.RefSchema:
 		return g.resolveRefSchema(s)
 	case *avro.RecordSchema:
-		return g.generate(s)
+		return g.resolveRecordSchema(s)
 	case *avro.PrimitiveSchema:
 		typ := primitiveMappings[s.Type()]
 		if ls := s.Logical(); ls != nil {
@@ -182,7 +174,7 @@ func (g *Generator) resolveType(schema avro.Schema) string {
 		}
 		return typ
 	case *avro.ArraySchema:
-		return fmt.Sprintf("[]%s", g.generate(s.Items()))
+		return "[]" + g.generate(s.Items())
 	case *avro.EnumSchema:
 		return "string"
 	case *avro.FixedSchema:
@@ -192,7 +184,7 @@ func (g *Generator) resolveType(schema avro.Schema) string {
 		}
 		return typ
 	case *avro.MapSchema:
-		return "map[string]" + g.resolveType(s.Values())
+		return "map[string]" + g.generate(s.Values())
 	case *avro.UnionSchema:
 		return g.resolveUnionTypes(s)
 	default:
@@ -207,16 +199,28 @@ func (g *Generator) resolveTypeName(s avro.NamedSchema) string {
 	return strcase.ToGoPascal(s.Name())
 }
 
+func (g *Generator) resolveRecordSchema(schema *avro.RecordSchema) string {
+	fields := make([]field, len(schema.Fields()))
+	for i, f := range schema.Fields() {
+		typ := g.generate(f.Type())
+		tag := f.Name()
+		fields[i] = g.newField(strcase.ToGoPascal(f.Name()), typ, tag)
+	}
+
+	typeName := g.resolveTypeName(schema)
+	g.typedefs = append(g.typedefs, newType(typeName, fields))
+	return typeName
+}
+
 func (g *Generator) resolveRefSchema(s *avro.RefSchema) string {
-	typ := ""
 	if sx, ok := s.Schema().(*avro.RecordSchema); ok {
-		typ = g.resolveTypeName(sx)
+		return g.resolveTypeName(sx)
 	}
-	return typ
+	return g.generate(s.Schema())
 }
 
 func (g *Generator) resolveUnionTypes(s *avro.UnionSchema) string {
-	types := make([]string, 0)
+	types := make([]string, 0, len(s.Types()))
 	for _, elem := range s.Types() {
 		if _, ok := elem.(*avro.NullSchema); ok {
 			continue


Test patch:
diff --git a/cmd/avrogen/testdata/golden.go b/cmd/avrogen/testdata/golden.go
index 569bc37..7f81cc5 100644
--- a/cmd/avrogen/testdata/golden.go
+++ b/cmd/avrogen/testdata/golden.go
@@ -1,6 +1,7 @@
-// Code generated by avro/gen. DO NOT EDIT.
 package testpkg
 
+// Code generated by avro/gen. DO NOT EDIT.
+
 type Test struct {
 	SomeString string `avro:"someString"`
 }
diff --git a/cmd/avrogen/testdata/golden_fullname.go b/cmd/avrogen/testdata/golden_fullname.go
index 697aca7..6434b86 100644
--- a/cmd/avrogen/testdata/golden_fullname.go
+++ b/cmd/avrogen/testdata/golden_fullname.go
@@ -1,6 +1,7 @@
-// Code generated by avro/gen. DO NOT EDIT.
 package testpkg
 
+// Code generated by avro/gen. DO NOT EDIT.
+
 type ABTest struct {
 	SomeString string `avro:"someString"`
 }
diff --git a/gen/gen_test.go b/gen/gen_test.go
index 9e2c7a3..b21d618 100644
--- a/gen/gen_test.go
+++ b/gen/gen_test.go
@@ -132,7 +132,7 @@ func TestStruct_GenFromRecordSchema(t *testing.T) {
 
 	want, err := os.ReadFile("testdata/golden.go")
 	require.NoError(t, err)
-	assert.Equal(t, want, file)
+	assert.Equal(t, string(want), string(file))
 }
 
 func TestStruct_GenFromRecordSchemaWithFullName(t *testing.T) {
@@ -149,7 +149,7 @@ func TestStruct_GenFromRecordSchemaWithFullName(t *testing.T) {
 
 	want, err := os.ReadFile("testdata/golden_fullname.go")
 	require.NoError(t, err)
-	assert.Equal(t, want, file)
+	assert.Equal(t, string(want), string(file))
 }
 
 func TestGenerator(t *testing.T) {
diff --git a/gen/testdata/golden.go b/gen/testdata/golden.go
index 5997dfe..0dae757 100644
--- a/gen/testdata/golden.go
+++ b/gen/testdata/golden.go
@@ -1,6 +1,7 @@
-// Code generated by avro/gen. DO NOT EDIT.
 package something
 
+// Code generated by avro/gen. DO NOT EDIT.
+
 import (
 	"math/big"
 	"time"
diff --git a/gen/testdata/golden_fullname.go b/gen/testdata/golden_fullname.go
index 56b00b5..54f509f 100644
--- a/gen/testdata/golden_fullname.go
+++ b/gen/testdata/golden_fullname.go
@@ -1,6 +1,7 @@
-// Code generated by avro/gen. DO NOT EDIT.
 package something
 
+// Code generated by avro/gen. DO NOT EDIT.
+
 import (
 	"math/big"
 	"time"
diff --git a/gen/testdata/golden_multiple.go b/gen/testdata/golden_multiple.go
index a1688dc..9c8f31d 100644
--- a/gen/testdata/golden_multiple.go
+++ b/gen/testdata/golden_multiple.go
@@ -1,6 +1,7 @@
-// Code generated by avro/gen. DO NOT EDIT.
 package something
 
+// Code generated by avro/gen. DO NOT EDIT.
+
 type TestUnionType struct {
 	Field1 int64 `avro:"Field1"`
 	Field2 int   `avro:"Field2"`