@@ -24,9 +24,10 @@ import (
2424)
2525
2626var (
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
3233var (
@@ -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 , `
398400var %s = []int32{%s}
399401
400402func (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