add support for non-scalar inputs

generate/generate.go

@@ -73,7 +73,7 @@ func (g *generator) getArgument(arg *ast.VariableDefinition) (argument, error) {
 	graphQLName := arg.Variable
 	firstRest := strings.SplitN(graphQLName, "", 2)
 	goName := strings.ToLower(firstRest[0]) + firstRest[1]
-	goType, err := g.addTypeForInputType(arg.Type)
+	goType, err := g.getTypeForInputType(arg.Type)
 	if err != nil {
 		return argument{}, err
 	}
@@ -123,7 +123,7 @@ func (g *generator) addOperation(op *ast.OperationDefinition) error {
 		}
 	}
 
-	responseName, err := g.addTypeForOperation(op)
+	responseName, err := g.getTypeForOperation(op)
 	if err != nil {
 		return err
 	}

generate/types.go

@@ -25,58 +25,131 @@ func (g *generator) baseTypeForOperation(operation ast.Operation) *ast.Definitio
 	}
 }
 
-func (g *generator) addTypeForOperation(operation *ast.OperationDefinition) (name string, err error) {
+func (g *generator) getTypeForOperation(operation *ast.OperationDefinition) (name string, err error) {
 	// TODO: configure ResponseName format
 	name = operation.Name + "Response"
 
 	if def, ok := g.typeMap[name]; ok {
-		// TODO: if the name is taken, maybe try to find another?
+		// TODO: check for and handle conflicts a better way
-		return "", fmt.Errorf("%s already defined:\n%s", name, def)
+		return name, fmt.Errorf("%s already defined:\n%s", name, def)
 	}
 
-	builder := &typeBuilder{generator: g}
+	selectionSet, err := selections(operation.SelectionSet)
-	fmt.Fprintf(builder, "type %s ", name)
-	err = builder.writeTypedef(
-		g.baseTypeForOperation(operation.Operation), operation.SelectionSet)
 	if err != nil {
-		return "", err
+		return name, err
 	}
 
-	def := builder.String()
+	err = g.addTypeForDefinition(
-	g.typeMap[name] = def
+		name, g.baseTypeForOperation(operation.Operation), selectionSet)
-	return name, nil
+
+	return name, err
 }
 
-func (g *generator) addTypeForInputType(typ *ast.Type) (string, error) {
+func (g *generator) addTypeForDefinition(name string, typ *ast.Definition, selectionSet []selection) error {
 	builder := &typeBuilder{generator: g}
+	fmt.Fprintf(builder, "type %s ", name)
+	err := builder.writeTypedef(typ, selectionSet)
+	if err != nil {
+		return err
+	}
 
-	// TODO: handle non-scalar types (by passing ...something... as the
+	g.typeMap[name] = builder.String()
-	// SelectionSet?)
+	return nil
-	err := builder.writeType(typ, nil)
+}
 
+func (g *generator) getTypeForInputType(typ *ast.Type) (string, error) {
+	builder := &typeBuilder{generator: g}
+	err := builder.writeType(typ, selectionsForType(g, typ), false)
 	return builder.String(), err
 }
 
-func (builder *typeBuilder) writeField(field *ast.Field) error {
+type selection interface {
+	Alias() string
+	Name() string
+	Type() *ast.Type
+	SelectionSet() ([]selection, error)
+}
+
+type field struct{ field *ast.Field }
+
+func (s field) Alias() string { return s.field.Alias }
+func (s field) Name() string  { return s.field.Name }
+
+func (s field) Type() *ast.Type {
+	if s.field.Definition == nil {
+		return nil
+	}
+	return s.field.Definition.Type
+}
+
+func (s field) SelectionSet() ([]selection, error) {
+	return selections(s.field.SelectionSet)
+}
+
+func selections(selectionSet ast.SelectionSet) ([]selection, error) {
+	retval := make([]selection, len(selectionSet))
+	for i, selection := range selectionSet {
+		switch selection := selection.(type) {
+		case *ast.Field:
+			retval[i] = field{selection}
+		case *ast.FragmentSpread, *ast.InlineFragment:
+			return nil, fmt.Errorf("not implemented: %T", selection)
+		default:
+			return nil, fmt.Errorf("invalid selection type: %v", selection)
+		}
+	}
+	return retval, nil
+}
+
+type inputField struct {
+	*generator
+	field *ast.FieldDefinition
+}
+
+func (s inputField) Alias() string   { return s.field.Name }
+func (s inputField) Name() string    { return s.field.Name }
+func (s inputField) Type() *ast.Type { return s.field.Type }
+
+func (s inputField) SelectionSet() ([]selection, error) {
+	return selectionsForType(s.generator, s.field.Type), nil
+}
+
+func selectionsForType(g *generator, typ *ast.Type) []selection {
+	def := g.schema.Types[typ.Name()]
+	selectionSet := make([]selection, len(def.Fields))
+	for i, field := range def.Fields {
+		selectionSet[i] = inputField{g, field}
+	}
+	return selectionSet
+}
+
+func (builder *typeBuilder) writeField(selection selection) error {
 	var jsonName string
-	if field.Alias != "" {
+	if selection.Alias() != "" {
-		jsonName = field.Alias
+		jsonName = selection.Alias()
 	} else {
 		// TODO: is this case needed? tests don't seem to get here.
-		jsonName = field.Name
+		jsonName = selection.Name()
 	}
 	// We need an exportable name for JSON-marshaling.
-	goName := strings.Title(jsonName)
+	goName := upperFirst(jsonName)
 
 	builder.WriteString(goName)
 	builder.WriteRune(' ')
 
-	if field.Definition == nil {
+	typ := selection.Type()
+	if typ == nil {
 		// Unclear why gqlparser hasn't already rejected this,
 		// but empirically it might not.
-		return fmt.Errorf("undefined field %v", field)
+		return fmt.Errorf("undefined field %v", selection.Name())
 	}
-	err := builder.writeType(field.Definition.Type, field.SelectionSet)
+
+	selectionSet, err := selection.SelectionSet()
+	if err != nil {
+		return err
+	}
+
+	err = builder.writeType(typ, selectionSet, true)
 	if err != nil {
 		return err
 	}
@@ -90,7 +163,7 @@ func (builder *typeBuilder) writeField(field *ast.Field) error {
 	return nil
 }
 
-var graphQLNameToGoName = map[string]string{
+var builtinTypes = map[string]string{
 	"Int":     "int", // TODO: technically int32 is always enough, use that?
 	"Float":   "float64",
 	"String":  "string",
@@ -98,7 +171,7 @@ var graphQLNameToGoName = map[string]string{
 	"ID":      "string", // TODO: named type for IDs?
 }
 
-func (builder *typeBuilder) writeType(typ *ast.Type, selectionSet ast.SelectionSet) error {
+func (builder *typeBuilder) writeType(typ *ast.Type, selectionSet []selection, inline bool) error {
 	// gqlgen does slightly different things here since it defines names for
 	// all the intermediate types, but its implementation may be useful to crib
 	// from:
@@ -113,27 +186,36 @@ func (builder *typeBuilder) writeType(typ *ast.Type, selectionSet ast.SelectionS
 		builder.WriteString("*")
 	}
 
-	return builder.writeTypedef(builder.schema.Types[typ.Name()], selectionSet)
+	_, ok := builtinTypes[typ.Name()]
+	def := builder.schema.Types[typ.Name()]
+	if ok || inline {
+		return builder.writeTypedef(def, selectionSet)
+	}
+
+	// TODO: casing should be configurable?
+	name := lowerFirst(typ.Name())
+	builder.WriteString(name)
+	if _, ok := builder.typeMap[name]; ok {
+		return nil
+	}
+	// Writes a typedef elsewhere
+	return builder.addTypeForDefinition(name, def, selectionSet)
 }
 
-func (builder *typeBuilder) writeTypedef(typedef *ast.Definition, selectionSet ast.SelectionSet) error {
+func (builder *typeBuilder) writeTypedef(typedef *ast.Definition, selectionSet []selection) error {
 	switch typedef.Kind {
 	case ast.Object, ast.InputObject:
 		builder.WriteString("struct {\n")
-		for _, selection := range selectionSet {
+		for _, field := range selectionSet {
-			switch selection := selection.(type) {
+			err := builder.writeField(field)
-			case *ast.Field:
+			if err != nil {
-				builder.writeField(selection)
+				return err
-			case *ast.FragmentSpread, *ast.InlineFragment:
-				return fmt.Errorf("not implemented: %T", selection)
-			default:
-				return fmt.Errorf("invalid selection type: %v", selection)
 			}
 		}
 		builder.WriteString("}")
 		return nil
 	case ast.Scalar, ast.Enum:
-		goName := graphQLNameToGoName[typedef.Name]
+		goName := builtinTypes[typedef.Name]
 		// TODO(benkraft): Handle custom scalars and enums.
 		if goName == "" {
 			return fmt.Errorf("unknown scalar name: %s", typedef.Name)

generate/types_test.go

@@ -1,7 +1,9 @@
 package generate
 
 import (
+	"fmt"
 	"go/format"
+	"strings"
 	"testing"
 
 	"github.com/vektah/gqlparser"
@@ -16,31 +18,34 @@ func gofmt(src string) (string, error) {
 	return string(formatted), nil
 }
 
-func TestTypeForOperation(t *testing.T) {
+var schemaText = `
-	schema, err := gqlparser.LoadSchema(&ast.Source{Name: "test schema", Input: `
+	input UserQueryInput {
-		type AuthMethod {
+		email: String
-			provider: String
+		name: String
-			email: String
+		id: ID
-		}
-
-		type User {
-			id: ID!
-			name: String
-			emails: [String!]!
-			emailsOrNull: [String!]
-			emailsWithNulls: [String]!
-			emailsWithNullsOrNull: [String]
-			authMethods: [AuthMethod!]!
-		}
-
-		type Query {
-			user: User
-		}
-	`})
-	if err != nil {
-		t.Fatal(err)
 	}
 
+	type AuthMethod {
+		provider: String
+		email: String
+	}
+
+	type User {
+		id: ID!
+		name: String
+		emails: [String!]!
+		emailsOrNull: [String!]
+		emailsWithNulls: [String]!
+		emailsWithNullsOrNull: [String]
+		authMethods: [AuthMethod!]!
+	}
+
+	type Query {
+		user: User
+	}
+`
+
+func TestTypeForOperation(t *testing.T) {
 	tests := []struct {
 		name           string
 		operation      string
@@ -109,17 +114,23 @@ func TestTypeForOperation(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			queryDoc, graphqlError := gqlparser.LoadQuery(schema, test.operation)
+			schema, graphqlError := gqlparser.LoadSchema(
+				&ast.Source{Name: "test schema", Input: schemaText})
 			if graphqlError != nil {
 				t.Fatal(graphqlError)
 			}
 
+			queryDoc, graphqlListError := gqlparser.LoadQuery(schema, test.operation)
+			if graphqlListError != nil {
+				t.Fatal(graphqlListError)
+			}
+
 			if len(queryDoc.Operations) != 1 {
 				t.Fatalf("got %v operations, want 1", len(queryDoc.Operations))
 			}
 
 			g := newGenerator("test_package", schema)
-			name, err := g.addTypeForOperation(queryDoc.Operations[0])
+			name, err := g.getTypeForOperation(queryDoc.Operations[0])
 			if err != nil {
 				t.Error(err)
 			}
@@ -136,3 +147,80 @@ func TestTypeForOperation(t *testing.T) {
 		})
 	}
 }
+
+func TestTypeForInputType(t *testing.T) {
+	tests := []struct {
+		name           string
+		graphQLType    string
+		expectedGoType string
+		otherTypes     []string
+	}{{
+		`RequiredBuiltin`,
+		`String!`,
+		`string`,
+		nil,
+	}, {
+		`ListOfBuiltin`,
+		`[String]`,
+		`[]*string`,
+		nil,
+	}, {
+		`DefinedType`,
+		`UserQueryInput`,
+		`*userQueryInput`,
+		[]string{`type userQueryInput struct {
+			Email    *string ` + "`json:\"email\"`" + `
+			Name     *string ` + "`json:\"name\"`" + `
+			Id       *string ` + "`json:\"id\"`" + `
+		}`},
+	}}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			expectedGoCode := fmt.Sprintf(
+				"type Input %s\n\n%s", test.expectedGoType,
+				strings.Join(test.otherTypes, "\n\n"))
+			expectedGoCode, err := gofmt(expectedGoCode)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			extraSchemaText := fmt.Sprintf(
+				"extend type Query { testQuery(var: %s): User }", test.graphQLType)
+			schema, graphqlError := gqlparser.LoadSchema(
+				&ast.Source{Name: "test schema", Input: schemaText},
+				&ast.Source{Name: "test schema extension", Input: extraSchemaText},
+			)
+			if graphqlError != nil {
+				t.Fatal(graphqlError)
+			}
+
+			operation := fmt.Sprintf(
+				"query($var: %s) { testQuery(var: $var) { id } }", test.graphQLType)
+			queryDoc, graphqlListError := gqlparser.LoadQuery(schema, operation)
+			if graphqlListError != nil {
+				t.Fatal(graphqlListError)
+			}
+
+			g := newGenerator("test_package", schema)
+			goType, err := g.getTypeForInputType(
+				queryDoc.Operations[0].VariableDefinitions[0].Type)
+			if err != nil {
+				t.Error(err)
+			}
+
+			goCode := fmt.Sprintf("type Input %s\n\n%s", goType, g.Types())
+
+			// gofmt before comparing.
+			goCode, err = gofmt(goCode)
+			if err != nil {
+				t.Error(err)
+			}
+
+			if goCode != expectedGoCode {
+				t.Errorf("got:\n%v\nwant:\n%v\n", goCode, expectedGoCode)
+			}
+		})
+	}
+}

generate/util.go

@@ -0,0 +1,28 @@
+package generate
+
+import (
+	"unicode"
+	"unicode/utf8"
+)
+
+func reverse(slice []string) {
+	for left, right := 0, len(slice)-1; left < right; left, right = left+1, right-1 {
+		slice[left], slice[right] = slice[right], slice[left]
+	}
+}
+
+func changeFirst(s string, f func(rune) rune) string {
+	c, n := utf8.DecodeRuneInString(s)
+	if c == utf8.RuneError { // empty or invalid
+		return s
+	}
+	return string(f(c)) + s[n:]
+}
+
+func lowerFirst(s string) string {
+	return changeFirst(s, unicode.ToLower)
+}
+
+func upperFirst(s string) string {
+	return changeFirst(s, unicode.ToUpper)
+}