Add ability to generate int enums

plugin/modelgen/models.go

@@ -48,6 +48,7 @@ type Enum struct {
 	Description string
 	Name        string
 	Values      []*EnumValue
+	IntValues   bool // Whether generated types should be integers
 }
 
 type EnumValue struct {
@@ -175,6 +176,7 @@ func (m *Plugin) MutateConfig(cfg *config.Config) error {
 			it := &Enum{
 				Name:        schemaType.Name,
 				Description: schemaType.Description,
+				IntValues:   schemaType.Directives.ForName("intEnum") != nil,
 			}
 
 			for _, v := range schemaType.EnumValues {

plugin/modelgen/models.gotpl

@@ -37,21 +37,35 @@
 
 {{ range $enum := .Enums }}
 	{{ with .Description }} {{.|prefixLines "// "}} {{end}}
-	type {{.Name|go }} string
+	type {{.Name|go }} {{if $enum.IntValues }}int{{else}}string{{end}}
 	const (
-	{{- range $value := .Values}}
+	{{- range $index, $value := .Values}}
 		{{- with .Description}}
 			{{.|prefixLines "// "}}
 		{{- end}}
-		{{ $enum.Name|go }}{{ .Name|go }} {{$enum.Name|go }} = {{.Name|quote}}
+		{{ $enum.Name|go }}{{ .Name|go }} {{$enum.Name|go }} = {{if $enum.IntValues }}{{$index}}{{else}}{{.Name|quote}}{{end}}
 	{{- end }}
 	)
 
-	var All{{.Name|go }} = []{{ .Name|go }}{
-	{{- range $value := .Values}}
-		{{$enum.Name|go }}{{ .Name|go }},
-	{{- end }}
-	}
+	{{if $enum.IntValues }}
+		var {{.Name|go }}Name = map[int]string{
+			{{- range $index, $value := .Values}}
+				{{$index}}: {{$value.Name|quote}},
+			{{- end }}
+		}
+
+		var {{.Name|go }}Value = map[string]int{
+			{{- range $index, $value := .Values}}
+				{{$value.Name|quote}}: {{$index}},
+			{{- end }}
+		}
+	{{else}}
+		var All{{.Name|go }} = []{{ .Name|go }}{
+		{{- range $value := .Values}}
+			{{$enum.Name|go }}{{ .Name|go }},
+		{{- end }}
+		}
+	{{end}}
 
 	func (e {{.Name|go }}) IsValid() bool {
 		switch e {
@@ -62,24 +76,39 @@
 	}
 
 	func (e {{.Name|go }}) String() string {
-		return string(e)
+		{{ if $enum.IntValues -}}
+			return {{.Name|go }}Name[int(e)]
+		{{ else -}}
+			return string(e)
+		{{ end -}}
 	}
 
 	func (e *{{.Name|go }}) UnmarshalGQL(v interface{}) error {
-		str, ok := v.(string)
-		if !ok {
-			return fmt.Errorf("enums must be strings")
-		}
+		{{ if $enum.IntValues -}}
+			value, ok := v.(int)
+			if !ok {
+				return fmt.Errorf("enums must be integers")
+			}
+		{{ else -}}
+			value, ok := v.(string)
+			if !ok {
+				return fmt.Errorf("enums must be strings")
+			}
+		{{ end -}}
 
-		*e = {{ .Name|go }}(str)
+		*e = {{ .Name|go }}(value)
 		if !e.IsValid() {
-			return fmt.Errorf("%s is not a valid {{ .Name }}", str)
+			return fmt.Errorf("%v is not a valid {{ .Name }}", value)
 		}
 		return nil
 	}
 
 	func (e {{.Name|go }}) MarshalGQL(w io.Writer) {
-		fmt.Fprint(w, strconv.Quote(e.String()))
+		{{ if $enum.IntValues -}}
+			fmt.Fprint(w, e)
+		{{ else -}}
+			fmt.Fprint(w, strconv.Quote(e.String()))
+		{{ end -}}
 	}
 
 {{- end }}

plugin/modelgen/models_test.go

@@ -24,6 +24,7 @@ func TestModelGeneration(t *testing.T) {
 	require.True(t, cfg.Models.UserDefined("MissingTypeNotNull"))
 	require.True(t, cfg.Models.UserDefined("MissingTypeNullable"))
 	require.True(t, cfg.Models.UserDefined("MissingEnum"))
+	require.True(t, cfg.Models.UserDefined("MissingIntEnum"))
 	require.True(t, cfg.Models.UserDefined("MissingUnion"))
 	require.True(t, cfg.Models.UserDefined("MissingInterface"))
 	require.True(t, cfg.Models.UserDefined("TypeWithDescription"))

plugin/modelgen/out/existing.go

@@ -21,6 +21,8 @@ type ExistingInput struct {
 
 type ExistingEnum string
 
+type ExistingIntEnum int
+
 type ExistingInterface interface {
 	IsExistingInterface()
 }

plugin/modelgen/testdata/gqlgen.yml

@@ -13,6 +13,8 @@ models:
     model: github.com/99designs/gqlgen/plugin/modelgen/out.ExistingInput
   ExistingEnum:
     model: github.com/99designs/gqlgen/plugin/modelgen/out.ExistingEnum
+  ExistingIntEnum:
+    model: github.com/99designs/gqlgen/plugin/modelgen/out.ExistingIntEnum
   ExistingInterface:
     model: github.com/99designs/gqlgen/plugin/modelgen/out.ExistingInterface
   ExistingUnion:

plugin/modelgen/testdata/schema.graphql

@@ -1,3 +1,5 @@
+directive @intEnum on ENUM
+
 type Query {
     thisShoudlntGetGenerated: Boolean
 }
@@ -36,6 +38,11 @@ enum MissingEnum {
     Goodbye
 }
 
+enum MissingIntEnum @intEnum {
+    Hello
+    Goodbye
+}
+
 interface MissingInterface {
     name: String
 }
@@ -59,6 +66,11 @@ enum ExistingEnum {
     Goodbye
 }
 
+enum ExistingIntEnum @intEnum {
+    Hello
+    Goodbye
+}
+
 interface ExistingInterface {
     name: String
 }