@@ -24,10 +24,11 @@ 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" )
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
3334var (
@@ -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 , `
434436var %s = []int32{%s}
435437
436438func (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