Repository: hamba/avro
Pull request #396 (2024-05-09 18:42:29)

feat: allow strict type generation with avrogen
This PR adds the support for generating go types with strict size types (i.e `int32` instead of `int`).


Patch:
diff --git a/cmd/avrogen/main.go b/cmd/avrogen/main.go
index fea65f0..e2a70d5 100644
--- a/cmd/avrogen/main.go
+++ b/cmd/avrogen/main.go
@@ -23,6 +23,7 @@ type config struct {
 	Tags        string
 	FullName    bool
 	Encoders    bool
+	StrictTypes bool
 	Initialisms string
 }
 
@@ -39,8 +40,9 @@ func realMain(args []string, stdout, stderr io.Writer) int {
 	flgs.StringVar(&cfg.Tags, "tags", "", "The additional field tags <tag-name>:{snake|camel|upper-camel|kebab}>[,...]")
 	flgs.BoolVar(&cfg.FullName, "fullname", false, "Use the full name of the Record schema to create the struct name.")
 	flgs.BoolVar(&cfg.Encoders, "encoders", false, "Generate encoders for the structs.")
+	flgs.BoolVar(&cfg.StrictTypes, "strict-types", false, "Use strict type sizes (e.g. int32) during generation.")
 	flgs.StringVar(&cfg.Initialisms, "initialisms", "", "Custom initialisms <VAL>[,...] for struct and field names.")
-	flgs.StringVar(&cfg.TemplateFileName, "templateFileName", "", "Override output template with one loaded from file.")
+	flgs.StringVar(&cfg.TemplateFileName, "template-filename", "", "Override output template with one loaded from file.")
 	flgs.Usage = func() {
 		_, _ = fmt.Fprintln(stderr, "Usage: avrogen [options] schemas")
 		_, _ = fmt.Fprintln(stderr, "Options:")
@@ -78,6 +80,7 @@ func realMain(args []string, stdout, stderr io.Writer) int {
 		gen.WithEncoders(cfg.Encoders),
 		gen.WithInitialisms(initialisms),
 		gen.WithTemplate(string(template)),
+		gen.WithStrictTypes(cfg.StrictTypes),
 	}
 	g := gen.NewGenerator(cfg.Pkg, tags, opts...)
 	for _, file := range flgs.Args() {
diff --git a/gen/gen.go b/gen/gen.go
index 3bc18dc..720d620 100644
--- a/gen/gen.go
+++ b/gen/gen.go
@@ -22,6 +22,7 @@ type Config struct {
 	Tags        map[string]TagStyle
 	FullName    bool
 	Encoders    bool
+	StrictTypes bool
 	Initialisms []string
 }
 
@@ -44,15 +45,20 @@ const (
 //go:embed output_template.tmpl
 var outputTemplate string
 
-var primitiveMappings = map[avro.Type]string{
-	"string":  "string",
-	"bytes":   "[]byte",
-	"int":     "int",
-	"long":    "int64",
-	"float":   "float32",
-	"double":  "float64",
-	"boolean": "bool",
-}
+var (
+	primitiveMappings = map[avro.Type]string{
+		"string":  "string",
+		"bytes":   "[]byte",
+		"int":     "int",
+		"long":    "int64",
+		"float":   "float32",
+		"double":  "float64",
+		"boolean": "bool",
+	}
+	strictTypeMappings = map[string]string{
+		"int": "int32",
+	}
+)
 
 // Struct generates Go structs based on the schema and writes them to w.
 func Struct(s string, w io.Writer, cfg Config) error {
@@ -74,6 +80,7 @@ func StructFromSchema(schema avro.Schema, w io.Writer, cfg Config) error {
 		WithFullName(cfg.FullName),
 		WithEncoders(cfg.Encoders),
 		WithInitialisms(cfg.Initialisms),
+		WithStrictTypes(cfg.StrictTypes),
 	}
 	g := NewGenerator(strcase.ToSnake(cfg.PackageName), cfg.Tags, opts...)
 	g.Parse(rec)
@@ -133,6 +140,13 @@ func WithTemplate(template string) OptsFunc {
 	}
 }
 
+// WithStrictTypes configures the generator to use strict type sizes.
+func WithStrictTypes(b bool) OptsFunc {
+	return func(g *Generator) {
+		g.strictTypes = b
+	}
+}
+
 // Generator generates Go structs from schemas.
 type Generator struct {
 	template    string
@@ -140,6 +154,7 @@ type Generator struct {
 	tags        map[string]TagStyle
 	fullName    bool
 	encoders    bool
+	strictTypes bool
 	initialisms []string
 
 	imports           []string
@@ -201,6 +216,11 @@ func (g *Generator) generate(schema avro.Schema) string {
 		if ls := s.Logical(); ls != nil {
 			typ = g.resolveLogicalSchema(ls.Type())
 		}
+		if g.strictTypes {
+			if newTyp, ok := strictTypeMappings[typ]; ok {
+				typ = newTyp
+			}
+		}
 		return typ
 	case *avro.ArraySchema:
 		return "[]" + g.generate(s.Items())


Test patch:
diff --git a/cmd/avrogen/main_test.go b/cmd/avrogen/main_test.go
index dc6a5b6..31f089f 100644
--- a/cmd/avrogen/main_test.go
+++ b/cmd/avrogen/main_test.go
@@ -143,6 +143,29 @@ func TestAvroGen_GeneratesSchemaWithEncoders(t *testing.T) {
 	assert.Equal(t, want, got)
 }
 
+func TestAvroGen_GeneratesSchemaWithStrictTypes(t *testing.T) {
+	path, err := os.MkdirTemp("./", "avrogen")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = os.RemoveAll(path) })
+
+	file := filepath.Join(path, "test.go")
+	args := []string{"avrogen", "-pkg", "testpkg", "-o", file, "-strict-types", "testdata/schema.avsc"}
+	gotCode := realMain(args, io.Discard, io.Discard)
+	require.Equal(t, 0, gotCode)
+
+	got, err := os.ReadFile(file)
+	require.NoError(t, err)
+
+	if *update {
+		err = os.WriteFile("testdata/golden_stricttypes.go", got, 0600)
+		require.NoError(t, err)
+	}
+
+	want, err := os.ReadFile("testdata/golden_stricttypes.go")
+	require.NoError(t, err)
+	assert.Equal(t, want, got)
+}
+
 func TestParseTags(t *testing.T) {
 	tests := []struct {
 		name string
diff --git a/cmd/avrogen/testdata/golden.go b/cmd/avrogen/testdata/golden.go
index 2c00ae2..cf6eda1 100644
--- a/cmd/avrogen/testdata/golden.go
+++ b/cmd/avrogen/testdata/golden.go
@@ -5,4 +5,5 @@ package testpkg
 // Test is a generated struct.
 type Test struct {
 	SomeString string `avro:"someString"`
+	SomeInt    int    `avro:"someInt"`
 }
diff --git a/cmd/avrogen/testdata/golden_encoders.go b/cmd/avrogen/testdata/golden_encoders.go
index d071865..2cd1308 100644
--- a/cmd/avrogen/testdata/golden_encoders.go
+++ b/cmd/avrogen/testdata/golden_encoders.go
@@ -9,9 +9,10 @@ import (
 // Test is a generated struct.
 type Test struct {
 	SomeString string `avro:"someString"`
+	SomeInt    int    `avro:"someInt"`
 }
 
-var schemaTest = avro.MustParse(`{"name":"a.b.test","type":"record","fields":[{"name":"someString","type":"string"}]}`)
+var schemaTest = avro.MustParse(`{"name":"a.b.test","type":"record","fields":[{"name":"someString","type":"string"},{"name":"someInt","type":"int"}]}`)
 
 // Schema returns the schema for Test.
 func (o *Test) Schema() avro.Schema {
diff --git a/cmd/avrogen/testdata/golden_fullname.go b/cmd/avrogen/testdata/golden_fullname.go
index 2f175b6..43b51d3 100644
--- a/cmd/avrogen/testdata/golden_fullname.go
+++ b/cmd/avrogen/testdata/golden_fullname.go
@@ -5,4 +5,5 @@ package testpkg
 // ABTest is a generated struct.
 type ABTest struct {
 	SomeString string `avro:"someString"`
+	SomeInt    int    `avro:"someInt"`
 }
diff --git a/cmd/avrogen/testdata/golden_stricttypes.go b/cmd/avrogen/testdata/golden_stricttypes.go
new file mode 100644
index 0000000..49f3dcb
--- /dev/null
+++ b/cmd/avrogen/testdata/golden_stricttypes.go
@@ -0,0 +1,9 @@
+package testpkg
+
+// Code generated by avro/gen. DO NOT EDIT.
+
+// Test is a generated struct.
+type Test struct {
+	SomeString string `avro:"someString"`
+	SomeInt    int32  `avro:"someInt"`
+}
diff --git a/cmd/avrogen/testdata/schema.avsc b/cmd/avrogen/testdata/schema.avsc
index 6183305..8196629 100644
--- a/cmd/avrogen/testdata/schema.avsc
+++ b/cmd/avrogen/testdata/schema.avsc
@@ -3,6 +3,7 @@
   "name": "test",
   "namespace": "a.b",
   "fields": [
-    { "name": "someString", "type": "string" }
+    { "name": "someString", "type": "string" },
+    { "name": "someInt", "type": "int" }
   ]
 }
\ No newline at end of file
diff --git a/gen/gen_test.go b/gen/gen_test.go
index 6dd4411..786beed 100644
--- a/gen/gen_test.go
+++ b/gen/gen_test.go
@@ -93,6 +93,24 @@ func TestStruct_HandlesAdditionalInitialisms(t *testing.T) {
 	assert.Contains(t, lines, "type CIDOverHTTPRecord struct {")
 }
 
+func TestStruct_HandlesStrictTypes(t *testing.T) {
+	schema := `{
+  "type": "record",
+  "name": "test",
+  "fields": [
+    { "name": "someString", "type": "int" }
+  ]
+}`
+	gc := gen.Config{
+		PackageName: "Something",
+		StrictTypes: true,
+	}
+
+	_, lines := generate(t, schema, gc)
+
+	assert.Contains(t, lines, "SomeString int32 `avro:\"someString\"`")
+}
+
 func TestStruct_ConfigurableFieldTags(t *testing.T) {
 	schema := `{
   "type": "record",