add support for custom scalars -- mainly adding proper import machinery

README.md

@@ -112,7 +112,6 @@ Config options:
 - whether names should be exported
 - default handling for optional fields? (maybe generate a HasFoo, you can always ignore if you don't care)
 - generate mocks?
-- custom scalar types (or custom mappings for standard scalars, if you want a special ID type say)
 
 Runtime:
 - (+) basic tests for graphql package

example/caller.go

@@ -48,7 +48,7 @@ func Main() {
 		if err != nil {
 			return
 		}
-		fmt.Println("you are", viewerResp.Viewer.MyName)
+		fmt.Println("you are", viewerResp.Viewer.MyName, "created on", viewerResp.Viewer.CreatedAt.Format("2006-01-02"))
 
 	case 2:
 		username := os.Args[1]
@@ -56,7 +56,7 @@ func Main() {
 		if err != nil {
 			return
 		}
-		fmt.Println(username, "is", userResp.User.TheirName)
+		fmt.Println(username, "is", userResp.User.TheirName, "created on", userResp.User.CreatedAt.Format("2006-01-02"))
 
 	default:
 		err = fmt.Errorf("usage: %v [username]", os.Args[0])

example/generated.go

@@ -4,6 +4,7 @@ package example
 
 import (
 	"context"
+	"time"
 
 	"github.com/Khan/genqlient/graphql"
 )
@@ -13,7 +14,8 @@ type getUserResponse struct {
 }
 
 type getUserUser struct {
-	TheirName string `json:"theirName"`
+	TheirName string    `json:"theirName"`
+	CreatedAt time.Time `json:"createdAt"`
 }
 
 type getViewerResponse struct {
@@ -21,7 +23,8 @@ type getViewerResponse struct {
 }
 
 type getViewerViewerUser struct {
-	MyName string
+	MyName    string
+	CreatedAt time.Time `json:"createdAt"`
 }
 
 func getViewer(
@@ -36,6 +39,7 @@ func getViewer(
 query getViewer {
 	viewer {
 		MyName: name
+		createdAt
 	}
 }
 `,
@@ -63,6 +67,7 @@ func getUser(
 query getUser ($Login: String!) {
 	user(login: $Login) {
 		theirName: name
+		createdAt
 	}
 }
 `,

example/genqlient.graphql

@@ -1,6 +1,7 @@
 query getViewer {
   viewer {
     MyName: name
+    createdAt
   }
 }
 
@@ -8,5 +9,6 @@ query getViewer {
 query getUser($Login: String!) {
   user(login: $Login) {
     theirName: name
+    createdAt
   }
 }

example/genqlient.yaml

@@ -4,3 +4,8 @@ schema: schema.graphql
 queries:
 - genqlient.graphql
 generated: generated.go
+
+# We map github's DateTime type to Go's time.Time (which conveniently already
+# defines MarshalJSON and UnmarshalJSAON).
+scalars:
+  DateTime: time.Time

generate/config.go

@@ -5,7 +5,6 @@ import (
 	"go/token"
 	"io/ioutil"
 	"path/filepath"
-	"strings"
 
 	"gopkg.in/yaml.v2"
 )
@@ -71,24 +70,32 @@ type Config struct {
 	// TODO: what if you want to return err?
 	ClientGetter string `yaml:"client_getter"`
 
+	// A map from GraphQL scalar type name to Go fully-qualified type name for
+	// the types to use for any custom or builtin scalars.  By default, builtin
+	// scalars are mapped to the obvious Go types (String and ID to string, Int
+	// to int, Float to float64, and Boolean to bool), but this setting will
+	// extend or override those mappings.  These types must define MarshalJSON
+	// and UnmarshalJSON methods, or otherwise be convertible to JSON.
+	Scalars map[string]string `yaml:"scalars"`
+
 	// Set automatically to the filename of the config file itself.
 	configFilename string
 }
 
-// BaseDir returns the directory of the config-file (relative to which
+// baseDir returns the directory of the config-file (relative to which
 // all the other paths are resolved).
-func (c *Config) BaseDir() string {
+func (c *Config) baseDir() string {
 	return filepath.Dir(c.configFilename)
 }
 
 func (c *Config) ValidateAndFillDefaults(configFilename string) error {
 	c.configFilename = configFilename
 	// Make paths relative to config dir
-	c.Schema = filepath.Join(c.BaseDir(), c.Schema)
+	c.Schema = filepath.Join(c.baseDir(), c.Schema)
 	for i := range c.Operations {
-		c.Operations[i] = filepath.Join(c.BaseDir(), c.Operations[i])
+		c.Operations[i] = filepath.Join(c.baseDir(), c.Operations[i])
 	}
-	c.Generated = filepath.Join(c.BaseDir(), c.Generated)
+	c.Generated = filepath.Join(c.baseDir(), c.Generated)
 
 	if c.Package == "" {
 		abs, err := filepath.Abs(c.Generated)
@@ -107,24 +114,6 @@ func (c *Config) ValidateAndFillDefaults(configFilename string) error {
 	return nil
 }
 
-func (c *Config) ContextPackage() string {
-	if c.ContextType == "" {
-		return ""
-	}
-
-	i := strings.LastIndex(c.ContextType, ".")
-	return c.ContextType[:i]
-}
-
-func (c *Config) ContextTypeReference() string {
-	if c.ContextType == "" {
-		return ""
-	}
-
-	i := strings.LastIndex(c.ContextType, "/")
-	return c.ContextType[i+1:]
-}
-
 func ReadAndValidateConfig(filename string) (*Config, error) {
 	config := *defaultConfig
 	if filename != "" {

generate/example_test.go

@@ -59,7 +59,7 @@ func TestRunExample(t *testing.T) {
 	}
 
 	got := strings.TrimSpace(string(out))
-	want := "benjaminjkraft is Ben Kraft"
+	want := "benjaminjkraft is Ben Kraft created on 2009-08-03"
 	if got != want {
 		t.Errorf("output incorrect\ngot:\n%s\nwant:\n%s", got, want)
 	}

generate/generate.go

@@ -7,16 +7,16 @@ import (
 	"go/format"
 	"sort"
 	"strings"
+	"text/template"
 
 	"github.com/vektah/gqlparser/v2/ast"
 	"github.com/vektah/gqlparser/v2/formatter"
+	"golang.org/x/tools/imports"
 )
 
 // Set to true to test features that aren't yet really ready.
 var allowBrokenFeatures = false
 
-var fileTemplate = mustTemplate("operation.go.tmpl")
-
 // generator is the context for the codegen process (and ends up getting passed
 // to the template).
 type generator struct {
@@ -25,9 +25,14 @@ type generator struct {
 	// The list of operations for which to generate code.
 	Operations []operation
 	// The types needed for these operations.
-	typeMap    map[string]string
+	typeMap map[string]string
-	ImportJSON bool
+	// Imports needed for these operations, path -> alias and alias -> true
-	schema     *ast.Schema
+	imports     map[string]string
+	usedAliases map[string]bool
+	// Cache of loaded templates.
+	templateCache map[string]*template.Template
+	// Schema we are generating code against
+	schema *ast.Schema
 }
 
 // JSON tags in operation are for ExportOperations (see Config for details).
@@ -59,11 +64,29 @@ type argument struct {
 }
 
 func newGenerator(config *Config, schema *ast.Schema) *generator {
-	return &generator{
+	g := generator{
-		Config:  config,
+		Config:        config,
-		typeMap: map[string]string{},
+		typeMap:       map[string]string{},
-		schema:  schema,
+		imports:       map[string]string{},
+		usedAliases:   map[string]bool{},
+		templateCache: map[string]*template.Template{},
+		schema:        schema,
 	}
+
+	if g.Config.ClientGetter == "" {
+		_, err := g.addRef("github.com/Khan/genqlient/graphql.Client")
+		if err != nil {
+			panic(err)
+		}
+	}
+	if g.Config.ContextType != "" {
+		_, err := g.addRef(g.Config.ContextType)
+		if err != nil {
+			panic(err)
+		}
+	}
+
+	return &g
 }
 
 func (g *generator) Types() string {
@@ -163,7 +186,7 @@ func Generate(config *Config) (map[string][]byte, error) {
 		return nil, err
 	}
 
-	document, err := getAndValidateQueries(config.BaseDir(), config.Operations, schema)
+	document, err := getAndValidateQueries(config.baseDir(), config.Operations, schema)
 	if err != nil {
 		return nil, err
 	}
@@ -187,7 +210,7 @@ func Generate(config *Config) (map[string][]byte, error) {
 	}
 
 	var buf bytes.Buffer
-	err = fileTemplate.Execute(&buf, g)
+	err = g.execute("operation.go.tmpl", &buf, g)
 	if err != nil {
 		return nil, fmt.Errorf("could not render template: %v", err)
 	}
@@ -198,9 +221,14 @@ func Generate(config *Config) (map[string][]byte, error) {
 		return nil, fmt.Errorf("could not gofmt code: %v\n---unformatted code---\n%v",
 			err, string(unformatted))
 	}
+	importsed, err := imports.Process(config.Generated, formatted, nil)
+	if err != nil {
+		return nil, fmt.Errorf("could not goimports code: %v\n---unimportsed code---\n%v",
+			err, string(formatted))
+	}
 
 	retval := map[string][]byte{
-		config.Generated: formatted,
+		config.Generated: importsed,
 	}
 
 	if config.ExportOperations != "" {

generate/generate_test.go

@@ -71,6 +71,10 @@ func TestGenerate(t *testing.T) {
 				Package:          "test",
 				Generated:        goFilename,
 				ExportOperations: queriesFilename,
+				Scalars: map[string]string{
+					"ID":       "github.com/me/mypkg.ID",
+					"DateTime": "time.Time",
+				},
 			})
 			if err != nil {
 				t.Fatal(err)

generate/imports.go

@@ -0,0 +1,81 @@
+package generate
+
+import (
+	"fmt"
+	"go/types"
+	"strconv"
+	"strings"
+)
+
+func (g *generator) addImportFor(pkgPath string) (alias string) {
+	if alias, ok := g.imports[pkgPath]; ok {
+		return alias
+	}
+
+	pkgName := pkgPath[strings.LastIndex(pkgPath, "/")+1:]
+	alias = pkgName
+	suffix := 2
+	for g.usedAliases[alias] {
+		alias = pkgName + strconv.Itoa(suffix)
+	}
+
+	g.imports[pkgPath] = alias
+	g.usedAliases[alias] = true
+	return alias
+}
+
+// addRef adds any imports necessary to refer to the given name, and returns a
+// reference alias.Name for it.
+func (g *generator) addRef(fullyQualifiedName string) (qualifiedName string, err error) {
+	return g.getRef(fullyQualifiedName, true)
+}
+
+// ref returns a reference alias.Name for the given import, if its package was
+// already added (e.g. via addRef), and an error if not.
+func (g *generator) ref(fullyQualifiedName string) (qualifiedName string, err error) {
+	return g.getRef(fullyQualifiedName, false)
+}
+
+func (g *generator) getRef(fullyQualifiedName string, addImport bool) (qualifiedName string, err error) {
+	i := strings.LastIndex(fullyQualifiedName, ".")
+	if i == -1 {
+		if types.Universe.Lookup(fullyQualifiedName) == nil {
+			return "", fmt.Errorf(
+				`unknown name "%v"; expected a builtin or path/to/package.Name`, fullyQualifiedName)
+		}
+		return fullyQualifiedName, nil
+	}
+
+	pkgPath := fullyQualifiedName[:i]
+	localName := fullyQualifiedName[i+1:]
+	var alias string
+	if addImport {
+		alias = g.addImportFor(pkgPath)
+	} else {
+		var ok bool
+		alias, ok = g.imports[pkgPath]
+		if !ok {
+			return "", fmt.Errorf(`no alias defined for package "%v"`, pkgPath)
+		}
+	}
+	return alias + "." + localName, nil
+}
+
+// Returns the import-clause to use in the generated code.
+func (g *generator) Imports() string {
+	if len(g.imports) == 0 {
+		return ""
+	}
+
+	var builder strings.Builder
+	builder.WriteString("import (\n")
+	for path, alias := range g.imports {
+		if path == alias || strings.HasSuffix(path, "/"+alias) {
+			builder.WriteString("\t" + strconv.Quote(path) + "\n")
+		} else {
+			builder.WriteString("\t" + alias + " " + strconv.Quote(path) + "\n")
+		}
+	}
+	builder.WriteString(")\n\n")
+	return builder.String()
+}

generate/operation.go.tmpl

@@ -2,16 +2,7 @@ package {{.Config.Package}}
 
 // Code generated by github.com/Khan/genqlient, DO NOT EDIT.
 
-import (
+{{.Imports}}
-    {{if .Config.ContextType -}}
-    "{{.Config.ContextPackage}}"
-    {{end}}
-    {{- if .ImportJSON -}}
-    "encoding/json"
-    {{end}}
-
-    "github.com/Khan/genqlient/graphql"
-)
 
 {{/* TODO: type-assert that your ctx type implements context.Context */}}
 
@@ -21,10 +12,10 @@ import (
 {{.Doc}}
 func {{.Name}}(
     {{if $.Config.ContextType -}}
-    ctx {{$.Config.ContextTypeReference}},
+    ctx {{ref $.Config.ContextType}},
     {{end}}
     {{- if not $.Config.ClientGetter -}}
-    client graphql.Client,
+    client {{ref "github.com/Khan/genqlient/graphql.Client"}},
     {{end}}
     {{- range .Args -}}
     {{.GoName}} {{.GoType}},

generate/template.go

@@ -1,6 +1,8 @@
 package generate
 
 import (
+	"fmt"
+	"io"
 	"path/filepath"
 	"runtime"
 	"text/template"
@@ -13,6 +15,24 @@ var (
 	thisDir               = filepath.Dir(thisFilename)
 )
 
-func mustTemplate(relFilename string) *template.Template {
+// execute executes the given template with the funcs from this generator.
-	return template.Must(template.ParseFiles(filepath.Join(thisDir, relFilename)))
+func (g *generator) execute(tmplRelFilename string, w io.Writer, data interface{}) error {
+	tmpl := g.templateCache[tmplRelFilename]
+	if tmpl == nil {
+		absFilename := filepath.Join(thisDir, tmplRelFilename)
+		funcMap := template.FuncMap{
+			"ref": g.ref,
+		}
+		var err error
+		tmpl, err = template.New(tmplRelFilename).Funcs(funcMap).ParseFiles(absFilename)
+		if err != nil {
+			return fmt.Errorf("could not load template %v: %v", absFilename, err)
+		}
+		g.templateCache[tmplRelFilename] = tmpl
+	}
+	err := tmpl.Execute(w, data)
+	if err != nil {
+		return fmt.Errorf("could not render template: %v", err)
+	}
+	return nil
 }

generate/testdata/queries/DateTime.graphql

@@ -0,0 +1,3 @@
+query convertTimezone($dt: DateTime!, $tz: String) {
+  convert(dt: $dt, tz: $tz)
+}

generate/testdata/queries/DateTime.graphql.go

@@ -0,0 +1,38 @@
+package test
+
+// Code generated by github.com/Khan/genqlient, DO NOT EDIT.
+
+import (
+	"time"
+
+	"github.com/Khan/genqlient/graphql"
+)
+
+type convertTimezoneResponse struct {
+	Convert time.Time `json:"convert"`
+}
+
+func convertTimezone(
+	client graphql.Client,
+	dt time.Time,
+	tz string,
+) (*convertTimezoneResponse, error) {
+	variables := map[string]interface{}{
+		"dt": dt,
+		"tz": tz,
+	}
+
+	var retval convertTimezoneResponse
+	err := client.MakeRequest(
+		nil,
+		"convertTimezone",
+		`
+query convertTimezone ($dt: DateTime!, $tz: String) {
+	convert(dt: $dt, tz: $tz)
+}
+`,
+		&retval,
+		variables,
+	)
+	return &retval, err
+}

generate/testdata/queries/DateTime.graphql.json

@@ -0,0 +1,9 @@
+{
+  "operations": [
+    {
+      "operationName": "convertTimezone",
+      "query": "\nquery convertTimezone ($dt: DateTime!, $tz: String) {\n\tconvert(dt: $dt, tz: $tz)\n}\n",
+      "sourceLocation": "testdata/queries/DateTime.graphql"
+    }
+  ]
+}

generate/testdata/queries/InputObject.graphql.go

@@ -4,6 +4,7 @@ package test
 
 import (
 	"github.com/Khan/genqlient/graphql"
+	"github.com/me/mypkg"
 )
 
 type InputObjectQueryResponse struct {
@@ -11,7 +12,7 @@ type InputObjectQueryResponse struct {
 }
 
 type InputObjectQueryUser struct {
-	Id string `json:"id"`
+	Id mypkg.ID `json:"id"`
 }
 
 type Role string
@@ -24,7 +25,7 @@ const (
 type UserQueryInput struct {
 	Email string   `json:"email"`
 	Name  string   `json:"name"`
-	Id    string   `json:"id"`
+	Id    mypkg.ID `json:"id"`
 	Role  Role     `json:"role"`
 	Names []string `json:"names"`
 }

generate/testdata/queries/InterfaceNoFragments.graphql.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 
 	"github.com/Khan/genqlient/graphql"
+	"github.com/me/mypkg"
 )
 
 type InterfaceNoFragmentsQueryResponse struct {
@@ -13,7 +14,7 @@ type InterfaceNoFragmentsQueryResponse struct {
 }
 
 type InterfaceNoFragmentsQueryRootTopic struct {
-	Id       string                                              `json:"id"`
+	Id       mypkg.ID                                            `json:"id"`
 	Name     string                                              `json:"name"`
 	Children []InterfaceNoFragmentsQueryRootTopicChildrenContent `json:"-"`
 }
@@ -66,8 +67,8 @@ func (v *InterfaceNoFragmentsQueryRootTopic) UnmarshalJSON(b []byte) error {
 }
 
 type InterfaceNoFragmentsQueryRootTopicChildrenArticle struct {
-	Id   string `json:"id"`
+	Id   mypkg.ID `json:"id"`
-	Name string `json:"name"`
+	Name string   `json:"name"`
 }
 
 func (v InterfaceNoFragmentsQueryRootTopicChildrenArticle) implementsGraphQLInterfaceInterfaceNoFragmentsQueryRootTopicChildrenContent() {
@@ -78,16 +79,16 @@ type InterfaceNoFragmentsQueryRootTopicChildrenContent interface {
 }
 
 type InterfaceNoFragmentsQueryRootTopicChildrenTopic struct {
-	Id   string `json:"id"`
+	Id   mypkg.ID `json:"id"`
-	Name string `json:"name"`
+	Name string   `json:"name"`
 }
 
 func (v InterfaceNoFragmentsQueryRootTopicChildrenTopic) implementsGraphQLInterfaceInterfaceNoFragmentsQueryRootTopicChildrenContent() {
 }
 
 type InterfaceNoFragmentsQueryRootTopicChildrenVideo struct {
-	Id   string `json:"id"`
+	Id   mypkg.ID `json:"id"`
-	Name string `json:"name"`
+	Name string   `json:"name"`
 }
 
 func (v InterfaceNoFragmentsQueryRootTopicChildrenVideo) implementsGraphQLInterfaceInterfaceNoFragmentsQueryRootTopicChildrenContent() {

generate/testdata/queries/ListInput.graphql.go

@@ -4,6 +4,7 @@ package test
 
 import (
 	"github.com/Khan/genqlient/graphql"
+	"github.com/me/mypkg"
 )
 
 type ListInputQueryResponse struct {
@@ -11,7 +12,7 @@ type ListInputQueryResponse struct {
 }
 
 type ListInputQueryUser struct {
-	Id string `json:"id"`
+	Id mypkg.ID `json:"id"`
 }
 
 func ListInputQuery(

generate/testdata/queries/QueryWithAlias.graphql.go

@@ -4,6 +4,7 @@ package test
 
 import (
 	"github.com/Khan/genqlient/graphql"
+	"github.com/me/mypkg"
 )
 
 type QueryWithAliasResponse struct {
@@ -11,7 +12,7 @@ type QueryWithAliasResponse struct {
 }
 
 type QueryWithAliasUser struct {
-	ID string
+	ID mypkg.ID
 }
 
 func QueryWithAlias(

generate/testdata/queries/QueryWithDoubleAlias.graphql.go

@@ -4,6 +4,7 @@ package test
 
 import (
 	"github.com/Khan/genqlient/graphql"
+	"github.com/me/mypkg"
 )
 
 type QueryWithDoubleAliasResponse struct {
@@ -11,8 +12,8 @@ type QueryWithDoubleAliasResponse struct {
 }
 
 type QueryWithDoubleAliasUser struct {
-	ID     string
+	ID     mypkg.ID
-	AlsoID string
+	AlsoID mypkg.ID
 }
 
 func QueryWithDoubleAlias(

generate/testdata/queries/SimpleInput.graphql.go

@@ -4,6 +4,7 @@ package test
 
 import (
 	"github.com/Khan/genqlient/graphql"
+	"github.com/me/mypkg"
 )
 
 type SimpleInputQueryResponse struct {
@@ -11,7 +12,7 @@ type SimpleInputQueryResponse struct {
 }
 
 type SimpleInputQueryUser struct {
-	Id string `json:"id"`
+	Id mypkg.ID `json:"id"`
 }
 
 func SimpleInputQuery(

generate/testdata/queries/SimpleMutation.graphql.go

@@ -4,11 +4,12 @@ package test
 
 import (
 	"github.com/Khan/genqlient/graphql"
+	"github.com/me/mypkg"
 )
 
 type SimpleMutationCreateUser struct {
-	Id   string `json:"id"`
+	Id   mypkg.ID `json:"id"`
-	Name string `json:"name"`
+	Name string   `json:"name"`
 }
 
 type SimpleMutationResponse struct {

generate/testdata/queries/SimpleQuery.graphql.go

@@ -4,6 +4,7 @@ package test
 
 import (
 	"github.com/Khan/genqlient/graphql"
+	"github.com/me/mypkg"
 )
 
 type SimpleQueryResponse struct {
@@ -11,7 +12,7 @@ type SimpleQueryResponse struct {
 }
 
 type SimpleQueryUser struct {
-	Id string `json:"id"`
+	Id mypkg.ID `json:"id"`
 }
 
 func SimpleQuery(

generate/testdata/queries/TypeName.graphql.go

@@ -4,6 +4,7 @@ package test
 
 import (
 	"github.com/Khan/genqlient/graphql"
+	"github.com/me/mypkg"
 )
 
 type TypeNameQueryResponse struct {
@@ -11,8 +12,8 @@ type TypeNameQueryResponse struct {
 }
 
 type TypeNameQueryUser struct {
-	Typename string `json:"__typename"`
+	Typename string   `json:"__typename"`
-	Id       string `json:"id"`
+	Id       mypkg.ID `json:"id"`
 }
 
 func TypeNameQuery(

generate/testdata/queries/schema.graphql

@@ -1,3 +1,5 @@
+scalar DateTime
+
 enum Role {
   STUDENT
   TEACHER
@@ -60,6 +62,7 @@ type Query {
   user(query: UserQueryInput): User
   root: Topic!
   randomLeaf: LeafContent!
+  convert(dt: DateTime!, tz: String): DateTime!
 }
 
 type Mutation {

generate/testdata/queries/unexported.graphql.go

@@ -4,6 +4,7 @@ package test
 
 import (
 	"github.com/Khan/genqlient/graphql"
+	"github.com/me/mypkg"
 )
 
 type Role string
@@ -18,13 +19,13 @@ type unexportedResponse struct {
 }
 
 type unexportedUser struct {
-	Id string `json:"id"`
+	Id mypkg.ID `json:"id"`
 }
 
 type userQueryInput struct {
 	Email string   `json:"email"`
 	Name  string   `json:"name"`
-	Id    string   `json:"id"`
+	Id    mypkg.ID `json:"id"`
 	Role  Role     `json:"role"`
 	Names []string `json:"names"`
 }

generate/types.go

@@ -62,8 +62,12 @@ var builtinTypes = map[string]string{
 }
 
 func (g *generator) addTypeForDefinition(namePrefix, nameOverride string, typ *ast.Definition, fields []field) (name string, err error) {
-	// If this is a builtin type, just refer to it.
+	// If this is a builtin type or custom scalar, just refer to it.
-	goName, ok := builtinTypes[typ.Name]
+	goName, ok := g.Config.Scalars[typ.Name]
+	if ok {
+		return g.addRef(goName)
+	}
+	goName, ok = builtinTypes[typ.Name]
 	if ok {
 		return goName, nil
 	}
@@ -314,8 +318,7 @@ func (builder *typeBuilder) writeTypedef(typedef *ast.Definition, fields []field
 		builder.WriteString(")\n")
 		return nil
 	case ast.Scalar:
-		// TODO(benkraft): Handle custom scalars.
+		return fmt.Errorf("unknown scalar %v: please add it to genqlient.yaml", typedef.Name)
-		return fmt.Errorf("not implemented: %v", typedef.Kind)
 	default:
 		return fmt.Errorf("unexpected kind: %v", typedef.Kind)
 	}

generate/unmarshal.go

@@ -1,7 +1,5 @@
 package generate
 
-var unmarshalTemplate = mustTemplate("unmarshal.go.tmpl")
-
 type templateData struct {
 	// Go type to which the method will be added
 	Type string
@@ -47,7 +45,11 @@ func (builder *typeBuilder) maybeWriteUnmarshal(fields []field) error {
 		return nil
 	}
 
-	builder.ImportJSON = true
+	_, err := builder.addRef("encoding/json.Unmarshal")
+	if err != nil {
+		return err
+	}
+
 	builder.WriteString("\n\n")
-	return unmarshalTemplate.Execute(builder, data)
+	return builder.execute("unmarshal.go.tmpl", builder, data)
 }

generate/unmarshal.go.tmpl

@@ -2,19 +2,19 @@ func (v *{{.Type}}) UnmarshalJSON(b []byte) error {
 	var firstPass struct{
 		*{{.Type}}
         {{range .Fields -}}
-		{{.GoName}} json.RawMessage `json:"{{.JSONName}}"`
+		{{.GoName}} {{ref "encoding/json.RawMessage"}} `json:"{{.JSONName}}"`
         {{end}}
 	}
 	firstPass.{{.Type}} = v
 
-	err := json.Unmarshal(b, &firstPass)
+	err := {{ref "encoding/json.Unmarshal"}}(b, &firstPass)
     if err != nil {
         return err
     }
 
     {{range .Fields -}}
 	var tn struct { TypeName string `json:"__typename"` }
-	err = json.Unmarshal(firstPass.{{.GoName}}, &tn)
+	err = {{ref "encoding/json.Unmarshal"}}(firstPass.{{.GoName}}, &tn)
     if err != nil {
         return err
     }
@@ -24,7 +24,7 @@ func (v *{{.Type}}) UnmarshalJSON(b []byte) error {
 	case "{{.GraphQLName}}":
         {{/* TODO: handle repeated fields! */}}
 		v.{{$field.GoName}} = {{.GoName}}{}
-		err = json.Unmarshal(
+		err = {{ref "encoding/json.Unmarshal"}}(
 			firstPass.{{$field.GoName}}, &v.{{$field.GoName}})
     {{end}}
     {{end}}

go.mod

@@ -4,5 +4,6 @@ go 1.13
 
 require (
 	github.com/vektah/gqlparser/v2 v2.1.0
+	golang.org/x/tools v0.0.0-20190125232054-d66bd3c5d5a6
 	gopkg.in/yaml.v2 v2.2.4
 )

go.sum

@@ -20,6 +20,7 @@ github.com/stretchr/testify v1.4.0 h1:2E4SXV/wtOkTonXsotYi4li6zVWxYlZuYNCXe9XRJy
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/vektah/gqlparser/v2 v2.1.0 h1:uiKJ+T5HMGGQM2kRKQ8Pxw8+Zq9qhhZhz/lieYvCMns=
 github.com/vektah/gqlparser/v2 v2.1.0/go.mod h1:SyUiHgLATUR8BiYURfTirrTcGpcE+4XkV2se04Px1Ms=
+golang.org/x/tools v0.0.0-20190125232054-d66bd3c5d5a6 h1:iZgcI2DDp6zW5v9Z/5+f0NuqoxNdmzg4hivjk2WLXpY=
 golang.org/x/tools v0.0.0-20190125232054-d66bd3c5d5a6/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo=