Skip to content

Commit 2631e8f

Browse files
author
Matt Jones
committed
add a flag for Enum Generate methods
1 parent 1821c94 commit 2631e8f

File tree

1 file changed

+19
-12
lines changed

1 file changed

+19
-12
lines changed

generator/go.go

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,11 @@ import (
2424
)
2525

2626
var (
27-
flagGoBinarystring = flag.Bool("go.binarystring", false, "Always use string for binary instead of []byte")
28-
flagGoJSONEnumnum = flag.Bool("go.json.enumnum", false, "For JSON marshal enums by number instead of name")
29-
flagGoPointers = flag.Bool("go.pointers", false, "Make all fields pointers")
30-
flagGoImportPrefix = flag.String("go.importprefix", "", "Prefix for thrift-generated go package imports")
27+
flagGoBinarystring = flag.Bool("go.binarystring", false, "Always use string for binary instead of []byte")
28+
flagGoJSONEnumnum = flag.Bool("go.json.enumnum", false, "For JSON marshal enums by number instead of name")
29+
flagGoPointers = flag.Bool("go.pointers", false, "Make all fields pointers")
30+
flagGoImportPrefix = flag.String("go.importprefix", "", "Prefix for thrift-generated go package imports")
31+
flagGoGenerateMethods = flag.Bool("go.generate", false, "Add testing/quick compatible Generate methods to enum types")
3132
)
3233

3334
var (
@@ -423,21 +424,23 @@ func (e *%s) UnmarshalJSON(b []byte) error {
423424
}
424425
`, enumName, enumName, enumName, enumName)
425426

426-
valueStrings := make([]string, 0, len(enum.Values))
427-
for _, val := range enum.Values {
428-
valueStrings = append(valueStrings, strconv.FormatInt(int64(val.Value), 10))
429-
}
430-
sort.Strings(valueStrings)
431-
valueStringsName := strings.ToLower(enumName) + "Values"
427+
if *flagGoGenerateMethods {
428+
valueStrings := make([]string, 0, len(enum.Values))
429+
for _, val := range enum.Values {
430+
valueStrings = append(valueStrings, strconv.FormatInt(int64(val.Value), 10))
431+
}
432+
sort.Strings(valueStrings)
433+
valueStringsName := strings.ToLower(enumName) + "Values"
432434

433-
g.write(out, `
435+
g.write(out, `
434436
var %s = []int32{%s}
435437
436438
func (e *%s) Generate(rand *rand.Rand, size int) reflect.Value {
437439
v := %s(%s[rand.Intn(%d)])
438440
return reflect.ValueOf(&v)
439441
}
440442
`, valueStringsName, strings.Join(valueStrings, ", "), enumName, enumName, valueStringsName, len(valueNames))
443+
}
441444

442445
return nil
443446
}
@@ -637,7 +640,11 @@ func (g *GoGenerator) generateSingle(out io.Writer, thriftPath string, thrift *p
637640
// Imports
638641
imports := []string{"fmt"}
639642
if len(thrift.Enums) > 0 {
640-
imports = append(imports, "strconv", "math/rand", "reflect")
643+
imports = append(imports, "strconv")
644+
645+
if *flagGoGenerateMethods {
646+
imports = append(imports, "math/rand", "reflect")
647+
}
641648
}
642649
if len(thrift.Includes) > 0 {
643650
for _, path := range thrift.Includes {

0 commit comments

Comments
 (0)