feat: add custom initialisms support to gen package and avrogen (#303)

Author: hhromic

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
+}

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)
+		})
+	}
+}

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,17 +255,17 @@ 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 {
 	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)
+		fields[i] = g.newField(g.nameCaser.ToPascal(f.Name()), typ, tag)
 	}
 
 	typeName := g.resolveTypeName(schema)

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",