diff --git a/cmd/avrogen/main.go b/cmd/avrogen/main.go index 7ba35a71..243ef978 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 :{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 [,...] 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/cmd/avrogen/main_test.go b/cmd/avrogen/main_test.go index 42b73011..0ea70271 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.go b/gen/gen.go index 7deb9521..b57d4957 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) diff --git a/gen/gen_test.go b/gen/gen_test.go index 738eb118..23998809 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",