Skip to content

Commit 3f52a50

Browse files
author
Matt Jones
committed
add a flag for Enum Generate methods
1 parent f857951 commit 3f52a50

File tree

1 file changed

+18
-11
lines changed

1 file changed

+18
-11
lines changed

generator/go.go

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,10 @@ 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")
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+
flagGoGenerateMethods = flag.Bool("go.generate", false, "Add testing/quick compatible Generate methods to enum types")
3031
)
3132

3233
var (
@@ -387,21 +388,23 @@ func (e *%s) UnmarshalJSON(b []byte) error {
387388
}
388389
`, enumName, enumName, enumName, enumName)
389390

390-
valueStrings := make([]string, 0, len(enum.Values))
391-
for _, val := range enum.Values {
392-
valueStrings = append(valueStrings, strconv.FormatInt(int64(val.Value), 10))
393-
}
394-
sort.Strings(valueStrings)
395-
valueStringsName := strings.ToLower(enumName) + "Values"
391+
if *flagGoGenerateMethods {
392+
valueStrings := make([]string, 0, len(enum.Values))
393+
for _, val := range enum.Values {
394+
valueStrings = append(valueStrings, strconv.FormatInt(int64(val.Value), 10))
395+
}
396+
sort.Strings(valueStrings)
397+
valueStringsName := strings.ToLower(enumName) + "Values"
396398

397-
g.write(out, `
399+
g.write(out, `
398400
var %s = []int32{%s}
399401
400402
func (e *%s) Generate(rand *rand.Rand, size int) reflect.Value {
401403
v := %s(%s[rand.Intn(%d)])
402404
return reflect.ValueOf(&v)
403405
}
404406
`, valueStringsName, strings.Join(valueStrings, ", "), enumName, enumName, valueStringsName, len(valueNames))
407+
}
405408

406409
return nil
407410
}
@@ -601,7 +604,11 @@ func (g *GoGenerator) generateSingle(out io.Writer, thriftPath string, thrift *p
601604
// Imports
602605
imports := []string{"fmt"}
603606
if len(thrift.Enums) > 0 {
604-
imports = append(imports, "strconv", "math/rand", "reflect")
607+
imports = append(imports, "strconv")
608+
609+
if *flagGoGenerateMethods {
610+
imports = append(imports, "math/rand", "reflect")
611+
}
605612
}
606613
if len(thrift.Includes) > 0 {
607614
for _, path := range thrift.Includes {

0 commit comments

Comments
 (0)