Repository: hamba/avro
Pull request #303 (2023-09-05 22:49:46)

feat: add custom initialisms support to gen package and avrogen
Hello, this simple PR adds support for custom initialisms to the `gen` package and `avrogen` command.

Additional custom initialisms are necessary because `strcase.ToGoPascal` only uses GoLint's list of initialisms.
Currently, if an initialism in struct/field names is not part of Golint's, then `strcase` wrongly downcases it.

For example, consider the following schema containing the domain initialism `AID`:
```json
{
  "type": "record",
  "name": "AIDEventHTTPService",
  "namespace": "example",
  "fields" : [
    {"name": "AIDNumber", "type": "long"},
    {"name": "Desc", "type": "string"}
  ]
}

```

Currently, `avrogen` generates the following Go struct (struct/field names: `AID` -> `Aid`):
```go
package example

// Code generated by avro/gen. DO NOT EDIT.

// AidEventHTTPService is a generated struct.
type AidEventHTTPService struct {
        AidNumber int64  `avro:"AIDNumber"`
        Desc      string `avro:"Desc"`
}
```

With this PR, `-initialisms AID` can be given to `avrogen` to generate the correct struct/field names:
```go
package example

// Code generated by avro/gen. DO NOT EDIT.

// AIDEventHTTPService is a generated struct.
type AIDEventHTTPService struct {
        AIDNumber int64  `avro:"AIDNumber"`
        Desc      string `avro:"Desc"`
}
```

There is a related PR in the upstream `strcase` repository: <https://github.com/ettle/strcase/pull/13>
That PR allows to preserve any extra initialisms without having to define them in advance, however the maintainers seems to not be merging it any time soon. This PR allows to solve the issue in the mean time.


hhromic: > One nitpick about test error assertion, otherwise looks good.

Thanks for the feedback, indeed looks cleaner after your suggestion.

I merged `main` too into this branch. Thanks for being open to contributions! Appreciated!

Patch:
diff --git a/cmd/avrogen/main.go b/cmd/avrogen/main.go
index 7ba35a7..243ef97 100644
--- a/cmd/avrogen/main.go
+++ b/cmd/avrogen/main.go
@@ -15,11 +15,12 @@ import (
 )
 
 type config struct {
-	Pkg      string
-	Out      string
-	Tags     string
-	FullName bool
-	Encoders bool
+	Pkg         string
+	Out         string
+	Tags        string
+	FullName    bool
+	Encoders    bool
+	Initialisms string
 }
 
 func main() {
@@ -35,6 +36,7 @@ func realMain(args []string, out, dumpout 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.StringVar(&cfg.Initialisms, "initialisms", "", "Custom initialisms <VAL>[,...] for struct and field names.")
 	flgs.Usage = func() {
 		_, _ = fmt.Fprintln(out, "Usage: avrogen [options] schemas")
 		_, _ = fmt.Fprintln(out, "Options:")
@@ -48,15 +50,23 @@ func realMain(args []string, out, dumpout io.Writer) int {
 		_, _ = fmt.Fprintln(out, "Error: "+err.Error())
 		return 1
 	}
+
 	tags, err := parseTags(cfg.Tags)
 	if err != nil {
 		_, _ = fmt.Fprintln(out, "Error: "+err.Error())
 		return 1
 	}
 
+	initialisms, err := parseInitialisms(cfg.Initialisms)
+	if err != nil {
+		_, _ = fmt.Fprintln(out, "Error: "+err.Error())
+		return 1
+	}
+
 	opts := []gen.OptsFunc{
 		gen.WithFullName(cfg.FullName),
 		gen.WithEncoders(cfg.Encoders),
+		gen.WithInitialisms(initialisms),
 	}
 	g := gen.NewGenerator(cfg.Pkg, tags, opts...)
 	for _, file := range flgs.Args() {
@@ -143,3 +153,19 @@ func parseTags(raw string) (map[string]gen.TagStyle, error) {
 	}
 	return result, nil
 }
+
+func parseInitialisms(raw string) ([]string, error) {
+	if raw == "" {
+		return []string{}, nil
+	}
+
+	result := []string{}
+	for _, initialism := range strings.Split(raw, ",") {
+		if initialism != strings.ToUpper(initialism) {
+			return nil, fmt.Errorf("initialism %q must be fully in upper case", initialism)
+		}
+		result = append(result, initialism)
+	}
+
+	return result, nil
+}
diff --git a/gen/gen.go b/gen/gen.go
index 7deb952..b57d495 100644
--- a/gen/gen.go
+++ b/gen/gen.go
@@ -20,6 +20,7 @@ type Config struct {
 	Tags        map[string]TagStyle
 	FullName    bool
 	Encoders    bool
+	Initialisms []string
 }
 
 // TagStyle defines the styling for a tag.
@@ -120,6 +121,7 @@ func StructFromSchema(schema avro.Schema, w io.Writer, cfg Config) error {
 	opts := []OptsFunc{
 		WithFullName(cfg.FullName),
 		WithEncoders(cfg.Encoders),
+		WithInitialisms(cfg.Initialisms),
 	}
 	g := NewGenerator(strcase.ToSnake(cfg.PackageName), cfg.Tags, opts...)
 	g.Parse(rec)
@@ -160,16 +162,27 @@ func WithEncoders(b bool) OptsFunc {
 	}
 }
 
+// WithInitialisms configures the generator to use additional custom initialisms
+// when styling struct and field names.
+func WithInitialisms(ss []string) OptsFunc {
+	return func(g *Generator) {
+		g.initialisms = ss
+	}
+}
+
 // Generator generates Go structs from schemas.
 type Generator struct {
-	pkg      string
-	tags     map[string]TagStyle
-	fullName bool
-	encoders bool
+	pkg         string
+	tags        map[string]TagStyle
+	fullName    bool
+	encoders    bool
+	initialisms []string
 
 	imports           []string
 	thirdPartyImports []string
 	typedefs          []typedef
+
+	nameCaser *strcase.Caser
 }
 
 // NewGenerator returns a generator.
@@ -183,6 +196,17 @@ func NewGenerator(pkg string, tags map[string]TagStyle, opts ...OptsFunc) *Gener
 		opt(g)
 	}
 
+	initialisms := map[string]bool{}
+	for _, v := range g.initialisms {
+		initialisms[v] = true
+	}
+
+	g.nameCaser = strcase.NewCaser(
+		true, // use standard Golint's initialisms
+		initialisms,
+		nil, // use default word split function
+	)
+
 	return g
 }
 
@@ -231,9 +255,9 @@ func (g *Generator) generate(schema avro.Schema) string {
 
 func (g *Generator) resolveTypeName(s avro.NamedSchema) string {
 	if g.fullName {
-		return strcase.ToGoPascal(s.FullName())
+		return g.nameCaser.ToPascal(s.FullName())
 	}
-	return strcase.ToGoPascal(s.Name())
+	return g.nameCaser.ToPascal(s.Name())
 }
 
 func (g *Generator) resolveRecordSchema(schema *avro.RecordSchema) string {
@@ -241,7 +265,7 @@ func (g *Generator) resolveRecordSchema(schema *avro.RecordSchema) string {
 	for i, f := range schema.Fields() {
 		typ := g.generate(f.Type())
 		tag := f.Name()
-		fields[i] = g.newField(strcase.ToGoPascal(f.Name()), typ, tag)
+		fields[i] = g.newField(g.nameCaser.ToPascal(f.Name()), typ, tag)
 	}
 
 	typeName := g.resolveTypeName(schema)


Test patch:
diff --git a/cmd/avrogen/main_test.go b/cmd/avrogen/main_test.go
index 42b7301..0ea7027 100644
--- a/cmd/avrogen/main_test.go
+++ b/cmd/avrogen/main_test.go
@@ -180,3 +180,36 @@ func TestParseTags(t *testing.T) {
 		})
 	}
 }
+
+func TestParseInitialisms(t *testing.T) {
+	tests := []struct {
+		name        string
+		initialisms string
+		errFunc     assert.ErrorAssertionFunc
+	}{
+		{
+			name:        "single initialism",
+			initialisms: "ABC",
+			errFunc:     assert.NoError,
+		},
+		{
+			name:        "multiple initialisms",
+			initialisms: "ABC,DEF",
+			errFunc:     assert.NoError,
+		},
+		{
+			name:        "wrong initialism",
+			initialisms: "ABC,def,GHI",
+			errFunc:     assert.Error,
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			_, err := parseInitialisms(test.initialisms)
+
+			test.errFunc(t, err)
+		})
+	}
+}
diff --git a/gen/gen_test.go b/gen/gen_test.go
index 738eb11..2399880 100644
--- a/gen/gen_test.go
+++ b/gen/gen_test.go
@@ -75,6 +75,24 @@ func TestStruct_HandlesGoInitialisms(t *testing.T) {
 	assert.Contains(t, lines, "type HTTPRecord struct {")
 }
 
+func TestStruct_HandlesAdditionalInitialisms(t *testing.T) {
+	schema := `{
+  "type": "record",
+  "name": "CidOverHttpRecord",
+  "fields": [
+    { "name": "someString", "type": "string" }
+  ]
+}`
+	gc := gen.Config{
+		PackageName: "Something",
+		Initialisms: []string{"CID"},
+	}
+
+	_, lines := generate(t, schema, gc)
+
+	assert.Contains(t, lines, "type CIDOverHTTPRecord struct {")
+}
+
 func TestStruct_ConfigurableFieldTags(t *testing.T) {
 	schema := `{
   "type": "record",