Repository: Khan/genqlient
Pull request #101 (2021-09-16 19:04:47)

Refactor import-loading to simplify the type-generation code
## Summary:
genqlient has some code (`imports.go`) dedicated to tracking which
imports we need and avoiding conflicts, as well as converting a
(restricted) Go expression like `map[string]github.com/me/mypkg.MyType`
to an import (`github.com/me/mypkg`) and a type-reference
(`map[string]mypkg.MyType`) to be used in the context of that import,
and at least making some attempt to track conflicts.  (Right now the
conflict-avoidance is not very smart, and not very well tested, but it
comes up rarely anyway.)  Sadly, that code was a bit cumbersome to use,
because you had to first register the imports (typically from
`convert.go`), then use them (often from the template).

In this commit I refactor the order we write things in order to allow a
significant simplification of how we import; in particular we no longer
have to guess in advance what imports which template will need; it can
just do `{{ref <expr>}}` as before, and it just works.  To do this, I:
- changed the importer to have only one API, which adds the import if
  needed, and returns the reference either way
- added a check that we don't add imports after they're written
- reorganized the toplevel templates a bit to make sure that check never
  fires; we now generate all the types and operations, then write the
  imports and glue it all together
This removes a bunch of silly code, and should simplify the process of
adding custom (un)marshalers (#38).

While I was at it, I put the documentation of what expressions we
support in a more visible place, and added a type-assertion that your
custom context type implements context.Context (if applicable).

## Test plan:
make check



Patch:
diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md
new file mode 100644
index 0000000..5859e82
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/bug_report.md
@@ -0,0 +1,23 @@
+---
+name: Bug report
+about: Create a report to help us improve
+title: ''
+labels: bug
+assignees: ''
+
+---
+
+**Describe the bug**
+A clear and concise description of what the bug is.
+
+**To Reproduce**
+Steps to reproduce the behavior. Please include whether the problem happens at code-generation time or at runtime.
+
+**Expected behavior**
+A clear and concise description of what you expected to happen.
+
+**genqlient version**
+Include your genqlient version here.
+
+**Additional context**
+Add any other context about the problem here.
diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md
new file mode 100644
index 0000000..11fc491
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/feature_request.md
@@ -0,0 +1,20 @@
+---
+name: Feature request
+about: Suggest an idea for this project
+title: ''
+labels: enhancement
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
+
+**Describe the solution you'd like**
+A clear and concise description of what you want to happen.
+
+**Describe alternatives you've considered**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context**
+Add any other context or screenshots about the feature request here.
diff --git a/.github/ISSUE_TEMPLATE/new-issue.md b/.github/ISSUE_TEMPLATE/new-issue.md
deleted file mode 100644
index b35fcf0..0000000
--- a/.github/ISSUE_TEMPLATE/new-issue.md
+++ /dev/null
@@ -1,10 +0,0 @@
----
-name: New issue
-about: gqlgen is not ready for production use!
-title: ''
-labels: ''
-assignees: ''
-
----
-
-genqlient is not ready for production use!  While bug reports and feature requests are welcome, we do not guarantee any response.
diff --git a/.github/ISSUE_TEMPLATE/security.md b/.github/ISSUE_TEMPLATE/security.md
new file mode 100644
index 0000000..7359c6c
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/security.md
@@ -0,0 +1,12 @@
+---
+name: Security
+about: Security issues should be reported on HackerOne.
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**STOP!** Security issues should be reported to the [Khan Academy HackerOne program](https://hackerone.com/khanacademy).
+
+**DO NOT REPORT SECURITY ISSUES ON GITHUB.**
diff --git a/README.md b/README.md
index 8926093..3ea90f3 100644
--- a/README.md
+++ b/README.md
@@ -1,8 +1,11 @@
 <img width="100%" alt="generated graphql client ⇒ genqlient" src="docs/images/genqlient.svg">
 
-# genqlient: a truly type-safe Go GraphQL client
+[![Go Reference](https://pkg.go.dev/badge/github.com/Khan/genqlient.svg)](https://pkg.go.dev/github.com/Khan/genqlient)
+[![Test Status](https://github.com/Khan/genqlient/actions/workflows/go.yml/badge.svg)](https://github.com/Khan/genqlient/actions)
+[![Contributor Covenant](https://img.shields.io/badge/Contributor%20Covenant-2.1-4baaaa.svg)](docs/CODE_OF_CONDUCT.md)
+[![GoReportcard](https://goreportcard.com/badge/github.com/Khan/genqlient?status.svg)](https://goreportcard.com/report/github.com/Khan/genqlient)
 
-This is a proof-of-concept of using code-generation to create a truly type-safe GraphQL client in Go.  It is not yet ready for contributions (see [docs/CONTRIBUTING.md](docs/CONTRIBUTING.md)).
+# genqlient: a truly type-safe Go GraphQL client
 
 ## What is genqlient?
 
@@ -20,7 +23,7 @@ You can download and run genqlient the usual way: `go run github.com/Khan/genqli
 
 ## How can I help?
 
-genqlient is not yet ready for contributions, but once it is, you'll be able to learn more in the ([Contribution Guidelines](docs/CONTRIBUTING.md)), or file an issue [on GitHub](issues).
+genqlient welcomes contributions!  Check out the ([Contribution Guidelines](docs/CONTRIBUTING.md)), or file an issue [on GitHub](issues).
 
 ## Why another GraphQL client?
 
diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
new file mode 100644
index 0000000..19ddb43
--- /dev/null
+++ b/docs/CHANGELOG.md
@@ -0,0 +1,31 @@
+# Changelog
+
+<!--
+
+When releasing a new version:
+- change "next" to the new version, and delete any empty sections
+- copy the below template to be the new "next"
+
+## template
+
+### Breaking changes:
+
+### New features:
+
+### Bug fixes:
+
+-->
+
+## next
+
+<!-- Add new changes in this section! -->
+
+### Breaking changes:
+
+### New features:
+
+### Bug fixes:
+
+## v0.1.0
+
+First open-sourced version.
diff --git a/docs/CODE_OF_CONDUCT.md b/docs/CODE_OF_CONDUCT.md
new file mode 100644
index 0000000..94c64fe
--- /dev/null
+++ b/docs/CODE_OF_CONDUCT.md
@@ -0,0 +1,132 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, caste, color, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+  overall community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+  advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+  address, without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement at
+opensource[at]khanacademy[dot]org.
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior,  harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.1, available at
+[https://www.contributor-covenant.org/version/2/1/code_of_conduct.html][v2.1].
+
+Community Impact Guidelines were inspired by
+[Mozilla's code of conduct enforcement ladder][Mozilla CoC].
+
+For answers to common questions about this code of conduct, see the FAQ at
+[https://www.contributor-covenant.org/faq][FAQ]. Translations are available
+at [https://www.contributor-covenant.org/translations][translations].
+
+[homepage]: https://www.contributor-covenant.org
+[v2.1]: https://www.contributor-covenant.org/version/2/1/code_of_conduct.html
+[Mozilla CoC]: https://github.com/mozilla/diversity
+[FAQ]: https://www.contributor-covenant.org/faq
+[translations]: https://www.contributor-covenant.org/translations
diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md
index 06b76ab..ab92499 100644
--- a/docs/CONTRIBUTING.md
+++ b/docs/CONTRIBUTING.md
@@ -1,8 +1,34 @@
 # Contributing to genqlient
 
-genqlient is not yet accepting contributions.  The library is still in a very rough state, so while we intend to accept contributions in the future, we're not ready for them just yet.  Please contact us (via an issue or email) if you are interested in helping out, but please be aware the answer may be that we aren't ready for your help yet, even if such help will be greatly useful once we are!  We'll be ready for the rest of the world soon.
+genqlient welcomes contributions from the community!  Our [help wanted](https://github.com/Khan/genqlient/issues?q=is%3Aissue+is%3Aopen+label%3A%22help+wanted%22) label tags issues where community help is especially welcome, and [good first issue](https://github.com/Khan/genqlient/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22) lists issues that are likely most approachable for new contributors.
 
-Khan Academy is a non-profit organization with a mission to provide a free, world-class education to anyone, anywhere. If you're looking for other ways to help us, You can help us in that mission by [donating](https://khanacademy.org/donate) or looking at [career opportunities](https://khanacademy.org/careers).
+While we're on the subject, Khan Academy is a non-profit organization with a mission to provide a free, world-class education to anyone, anywhere. If you're looking for other ways to help us, you can help us in that mission by [donating](https://khanacademy.org/donate) or looking at [career opportunities](https://khanacademy.org/careers).
+
+## Formalities
+
+In order to contribute to genqlient, you must first sign the [Khan CLA](https://www.khanacademy.org/r/cla). genqlient is released under the [MIT License](../LICENSE). While contributing, you must abide by the [Code of Conduct](CODE_OF_CONDUCT.md).
+
+## Security issues
+
+Security issues should be reported to the [Khan Academy HackerOne program](https://hackerone.com/khanacademy).
+
+## Pull requests
+
+Pull requests should have:
+
+- a clear PR title and description, describing the changes
+- a link to the issue fixed, if applicable
+- test coverage, as appropriate (see [below](#tests))
+- documentation, for new features
+- changelog entries
+
+Pull requests will be squash-merged, so subsequent commit messages may be brief (e.g. "review comments").
+
+Large changes should typically be discussed on the issue tracker first, and should ideally be broken up into separate PRs, or failing that, several commits, for ease of reviewing.
+
+## Style
+
+Go style should generally follow the conventions of [Effective Go](https://golang.org/doc/effective_go), and should have no lint errors (`make lint` to check). TODOs follow the [Khan Academy style](https://github.com/Khan/style-guides#todosmessage), i.e. `// TODO(yourusername): Simplify once we drop support for Go 1.23.` In general, try to make your code match the style of the surrounding code.
 
 ## Tests
 
@@ -12,7 +38,8 @@ Notes for contributors:
 - Most of the tests are snapshot-based; see `generate/generate_test.go`.  All new code-generation logic should be snapshot-tested.  Some code additionally has standalone unit tests, when convenient.
 - Integration tests run against a gqlgen server in `internal/integration/integration_test.go`, and should cover everything that snapshot tests can't, including the GraphQL client code and JSON marshaling.
 - If `GITHUB_TOKEN` is available in the environment, it also checks that the example returns the expected output when run against the real API.  This is configured automatically in GitHub Actions, but you can also use a [personal access token](https://docs.github.com/en/github/authenticating-to-github/creating-a-personal-access-token) with no scopes.  There's no need for this to cover anything in particular; it's just to make sure the example in fact works.
+- Tests should use `testify/assert` and `testify/require` where convenient (when making many simple assertions).
 
-## Design
+## Finding your way around
 
-See [DESIGN.md](DESIGN.md) for documentation of major design decisions in this library.
+If you're new to genqlient, start out by reading the source of `generate.Generate`, whose comments describe most of the high-level operation of genqlient.  In general, the code is documented inline, often with an introductory comment at the top of the file.  See [DESIGN.md](DESIGN.md) for documentation of major design decisions, which is a good way to get a sense of why genqlient is structured the way it is.
diff --git a/docs/SECURITY.md b/docs/SECURITY.md
new file mode 100644
index 0000000..264d3cb
--- /dev/null
+++ b/docs/SECURITY.md
@@ -0,0 +1,3 @@
+# Security policy
+
+Security issues should be reported to the [Khan Academy HackerOne program](https://hackerone.com/khanacademy).
diff --git a/docs/genqlient.yaml b/docs/genqlient.yaml
index d1a12bf..8dc540b 100644
--- a/docs/genqlient.yaml
+++ b/docs/genqlient.yaml
@@ -93,6 +93,16 @@ bindings:
     #  time.Time
     #  map[string]interface{}
     #  github.com/you/yourpkg/subpkg.MyType
+    # Specifically, this can be any of the following expressions:
+    # - any named type (qualified by the full package path)
+    # - any predeclared basic type (string, int, etc.)
+    # - interface{}
+    # - for any allowed type T, *T, []T, [N]T, and map[string]T
+    # but can't be, for example:
+    # - an inline (unnamed) struct or interface type
+    # - a map whose key-type is not string
+    # - a nonstandard way of spelling those, (interface {/* hi */},
+    #   map[  string      ]T)
     type: time.Time
 
   # To bind an object type:
diff --git a/generate/convert.go b/generate/convert.go
index 609c993..06271e8 100644
--- a/generate/convert.go
+++ b/generate/convert.go
@@ -165,7 +165,7 @@ func (g *generator) convertType(
 	// bind GraphQL named types, at least for now.
 	localBinding := options.Bind
 	if localBinding != "" && localBinding != "-" {
-		goRef, err := g.addRef(localBinding)
+		goRef, err := g.ref(localBinding)
 		return &goOpaqueType{goRef, typ.Name()}, err
 	}
 
@@ -217,7 +217,7 @@ func (g *generator) convertDefinition(
 				return nil, err
 			}
 		}
-		goRef, err := g.addRef(globalBinding.Type)
+		goRef, err := g.ref(globalBinding.Type)
 		return &goOpaqueType{goRef, def.Name}, err
 	}
 	goBuiltinName, ok := builtinTypes[def.Name]
diff --git a/generate/generate.go b/generate/generate.go
index 2adbf3c..08f2ce5 100644
--- a/generate/generate.go
+++ b/generate/generate.go
@@ -7,8 +7,8 @@ package generate
 import (
 	"bytes"
 	"encoding/json"
-	"fmt"
 	"go/format"
+	"io"
 	"sort"
 	"strings"
 	"text/template"
@@ -25,12 +25,15 @@ type generator struct {
 	// The config for which we are generating code.
 	Config *Config
 	// The list of operations for which to generate code.
-	Operations []operation
+	Operations []*operation
 	// The types needed for these operations.
 	typeMap map[string]goType
 	// Imports needed for these operations, path -> alias and alias -> true
 	imports     map[string]string
 	usedAliases map[string]bool
+	// True if we've already written out the imports (in which case they can't
+	// be modified).
+	importsLocked bool
 	// Cache of loaded templates.
 	templateCache map[string]*template.Template
 	// Schema we are generating code against
@@ -58,10 +61,12 @@ type operation struct {
 	ResponseName string `json:"-"`
 	// The original filename from which we got this query.
 	SourceFilename string `json:"sourceLocation"`
+	// The config within which we are generating code.
+	Config *Config `json:"-"`
 }
 
 type exportedOperations struct {
-	Operations []operation `json:"operations"`
+	Operations []*operation `json:"operations"`
 }
 
 type argument struct {
@@ -76,7 +81,7 @@ func newGenerator(
 	config *Config,
 	schema *ast.Schema,
 	fragments ast.FragmentDefinitionList,
-) (*generator, error) {
+) *generator {
 	g := generator{
 		Config:        config,
 		typeMap:       map[string]goType{},
@@ -91,29 +96,10 @@ func newGenerator(
 		g.fragments[fragment.Name] = fragment
 	}
 
-	_, err := g.addRef("github.com/Khan/genqlient/graphql.Client")
-	if err != nil {
-		return nil, err
-	}
-
-	if g.Config.ClientGetter != "" {
-		_, err := g.addRef(g.Config.ClientGetter)
-		if err != nil {
-			return nil, fmt.Errorf("invalid client_getter: %w", err)
-		}
-	}
-
-	if g.Config.ContextType != "-" {
-		_, err := g.addRef(g.Config.ContextType)
-		if err != nil {
-			return nil, fmt.Errorf("invalid context_type: %w", err)
-		}
-	}
-
-	return &g, nil
+	return &g
 }
 
-func (g *generator) Types() (string, error) {
+func (g *generator) WriteTypes(w io.Writer) error {
 	names := make([]string, 0, len(g.typeMap))
 	for name := range g.typeMap {
 		names = append(names, name)
@@ -124,17 +110,19 @@ func (g *generator) Types() (string, error) {
 	// vaguely aligned to the structure of the queries.
 	sort.Strings(names)
 
-	defs := make([]string, 0, len(g.typeMap))
-	var builder strings.Builder
 	for _, name := range names {
-		builder.Reset()
-		err := g.typeMap[name].WriteDefinition(&builder, g)
+		err := g.typeMap[name].WriteDefinition(w, g)
+		if err != nil {
+			return err
+		}
+		// Make sure we have blank lines between types (and between the last
+		// type and the first operation)
+		_, err = io.WriteString(w, "\n\n")
 		if err != nil {
-			return "", err
+			return err
 		}
-		defs = append(defs, builder.String())
 	}
-	return strings.Join(defs, "\n\n"), nil
+	return nil
 }
 
 func (g *generator) getArgument(
@@ -317,15 +305,18 @@ func (g *generator) addOperation(op *ast.OperationDefinition) error {
 		sourceFilename = sourceFilename[:i]
 	}
 
-	g.Operations = append(g.Operations, operation{
+	g.Operations = append(g.Operations, &operation{
 		Type: op.Operation,
 		Name: op.Name,
 		Doc:  docComment,
-		// The newline just makes it format a little nicer.
+		// The newline just makes it format a little nicer.  We add it here
+		// rather than in the template so exported operations will match
+		// *exactly* what we send to the server.
 		Body:           "\n" + builder.String(),
 		Args:           args,
 		ResponseName:   responseType.Reference(),
 		SourceFilename: sourceFilename,
+		Config:         g.Config, // for the convenience of the template
 	})
 
 	return nil
@@ -362,24 +353,55 @@ func Generate(config *Config) (map[string][]byte, error) {
 
 	// Step 2: For each operation and fragment, convert it into data structures
 	// representing Go types (defined in types.go).  The bulk of this logic is
-	// in convert.go.
-	g, err := newGenerator(config, schema, document.Fragments)
+	// in convert.go, and it additionally updates g.typeMap to include all the
+	// types it needs.
+	g := newGenerator(config, schema, document.Fragments)
+	for _, op := range document.Operations {
+		if err = g.addOperation(op); err != nil {
+			return nil, err
+		}
+	}
+
+	// Step 3: Glue it all together!
+	//
+	// First, write the types (from g.typeMap) and operations to a temporary
+	// buffer, since they affect what imports we'll put in the header.
+	var bodyBuf bytes.Buffer
+	err = g.WriteTypes(&bodyBuf)
 	if err != nil {
 		return nil, err
 	}
-	for _, op := range document.Operations {
-		if err = g.addOperation(op); err != nil {
+	for _, operation := range g.Operations {
+		err = g.render("operation.go.tmpl", &bodyBuf, operation)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	// The header also needs to reference some context types, which it does
+	// after it writes the imports, so we need to preregister those imports.
+	if g.Config.ContextType != "-" {
+		_, err = g.ref("context.Context")
+		if err != nil {
 			return nil, err
 		}
+		if g.Config.ContextType != "context.Context" {
+			_, err = g.ref(g.Config.ContextType)
+			if err != nil {
+				return nil, err
+			}
+		}
 	}
 
-	// Step 3: Glue it all together!  Most of this is done inline in the
-	// template, but the call to g.Types() in the template calls out to
-	// types.go to actually generate the code for each type.
+	// Now really glue it all together, and format.
 	var buf bytes.Buffer
-	err = g.execute("operation.go.tmpl", &buf, g)
+	err = g.render("header.go.tmpl", &buf, g)
 	if err != nil {
-		return nil, errorf(nil, "could not render template: %v", err)
+		return nil, err
+	}
+	_, err = io.Copy(&buf, &bodyBuf)
+	if err != nil {
+		return nil, err
 	}
 
 	unformatted := buf.Bytes()
diff --git a/generate/header.go.tmpl b/generate/header.go.tmpl
new file mode 100644
index 0000000..0055504
--- /dev/null
+++ b/generate/header.go.tmpl
@@ -0,0 +1,10 @@
+package {{.Config.Package}}
+
+// Code generated by github.com/Khan/genqlient, DO NOT EDIT.
+
+{{.Imports}}
+
+{{if and (ne .Config.ContextType "-") (ne .Config.ContextType "context.Context") }}
+// Check that context_type from genqlient.yaml implements context.Context.
+var _ {{ref "context.Context"}} = ({{ref .Config.ContextType}})(nil)
+{{end}}
diff --git a/generate/imports.go b/generate/imports.go
index e7a6007..89e5182 100644
--- a/generate/imports.go
+++ b/generate/imports.go
@@ -9,15 +9,12 @@ import (
 )
 
 func (g *generator) addImportFor(pkgPath string) (alias string) {
-	if existingAlias, ok := g.imports[pkgPath]; ok {
-		return existingAlias
-	}
-
 	pkgName := pkgPath[strings.LastIndex(pkgPath, "/")+1:]
 	alias = pkgName
 	suffix := 2
 	for g.usedAliases[alias] {
 		alias = pkgName + strconv.Itoa(suffix)
+		suffix++
 	}
 
 	g.imports[pkgPath] = alias
@@ -25,44 +22,33 @@ func (g *generator) addImportFor(pkgPath string) (alias string) {
 	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)
-}
-
 var _sliceOrMapPrefixRegexp = regexp.MustCompile(`^(\*|\[\d*\]|map\[string\])*`)
 
-func (g *generator) getRef(fullyQualifiedName string, addImport bool) (qualifiedName string, err error) {
-	// Ideally, we want to allow a reference to basically an arbitrary symbol.
-	// But that's very hard, because it might be quite complicated, like
-	//	struct{ F []map[mypkg.K]otherpkg.V }
-	// Now in practice, using an unnamed struct is not a great idea, but we do
-	// want to allow as much as we can that encoding/json knows how to work
-	// with, since you would reasonably expect us to accept, say,
-	// map[string][]interface{}.  So we allow:
-	// - any named type (mypkg.T)
-	// - any predeclared basic type (string, int, etc.)
-	// - interface{}
-	// - for any allowed type T, *T, []T, [N]T, and map[string]T
-	// which effectively excludes:
-	// - unnamed struct types
-	// - map[K]V where K is a named type wrapping string
-	// - any nonstandard spelling of those (interface {/* hi */},
-	//	 map[  string      ]T)
-	// TODO: document that somewhere visible
-
+// ref takes a Go fully-qualified name, ensures that any necessary symbols are
+// imported, and returns an appropriate reference.
+//
+// Ideally, we want to allow a reference to basically an arbitrary symbol.
+// But that's very hard, because it might be quite complicated, like
+//	struct{ F []map[mypkg.K]otherpkg.V }
+// Now in practice, using an unnamed struct is not a great idea, but we do
+// want to allow as much as we can that encoding/json knows how to work
+// with, since you would reasonably expect us to accept, say,
+// map[string][]interface{}.  So we allow:
+// - any named type (mypkg.T)
+// - any predeclared basic type (string, int, etc.)
+// - interface{}
+// - for any allowed type T, *T, []T, [N]T, and map[string]T
+// which effectively excludes:
+// - unnamed struct types
+// - map[K]V where K is a named type wrapping string
+// - any nonstandard spelling of those (interface {/* hi */},
+//	 map[  string      ]T)
+// (This is documented in docs/genqlient.yaml)
+func (g *generator) ref(fullyQualifiedName string) (qualifiedName string, err error) {
 	errorMsg := `invalid type-name "%v" (%v); expected a builtin, ` +
 		`path/to/package.Name, interface{}, or a slice, map, or pointer of those`
 
 	if strings.Contains(fullyQualifiedName, " ") {
-		// TODO: pass in pos here and below
 		return "", errorf(nil, errorMsg, fullyQualifiedName, "contains spaces")
 	}
 
@@ -80,22 +66,20 @@ func (g *generator) getRef(fullyQualifiedName string, addImport bool) (qualified
 
 	pkgPath := nameToImport[:i]
 	localName := nameToImport[i+1:]
-	var alias string
-	if addImport {
-		alias = g.addImportFor(pkgPath)
-	} else {
-		var ok bool
-		alias, ok = g.imports[pkgPath]
-		if !ok {
-			// This is an internal error, not a user error.
-			return "", errorf(nil, `no alias defined for package "%v"`, pkgPath)
+	alias, ok := g.imports[pkgPath]
+	if !ok {
+		if g.importsLocked {
+			return "", errorf(nil,
+				`genqlient internal error: imports locked but package "%v" has not been imported`, pkgPath)
 		}
+		alias = g.addImportFor(pkgPath)
 	}
 	return prefix + alias + "." + localName, nil
 }
 
 // Returns the import-clause to use in the generated code.
 func (g *generator) Imports() string {
+	g.importsLocked = true
 	if len(g.imports) == 0 {
 		return ""
 	}
diff --git a/generate/operation.go.tmpl b/generate/operation.go.tmpl
index c5614b5..e5cadf6 100644
--- a/generate/operation.go.tmpl
+++ b/generate/operation.go.tmpl
@@ -1,20 +1,9 @@
-package {{.Config.Package}}
-
-// Code generated by github.com/Khan/genqlient, DO NOT EDIT.
-
-{{.Imports}}
-
-{{/* TODO: type-assert that your ctx type implements context.Context */}}
-
-{{.Types}}
-
-{{range .Operations}}
 {{.Doc}}
 func {{.Name}}(
-    {{if ne $.Config.ContextType "-" -}}
-    ctx {{ref $.Config.ContextType}},
+    {{if ne .Config.ContextType "-" -}}
+    ctx {{ref .Config.ContextType}},
     {{end}}
-    {{- if not $.Config.ClientGetter -}}
+    {{- if not .Config.ClientGetter -}}
     client {{ref "github.com/Khan/genqlient/graphql.Client"}},
     {{end}}
     {{- range .Args -}}
@@ -46,8 +35,8 @@ func {{.Name}}(
     {{end -}}
 
     var err error
-    {{if $.Config.ClientGetter -}}
-    client, err := {{ref $.Config.ClientGetter}}({{if ne $.Config.ContextType "-"}}ctx{{else}}{{end}})
+    {{if .Config.ClientGetter -}}
+    client, err := {{ref .Config.ClientGetter}}({{if ne .Config.ContextType "-"}}ctx{{else}}{{end}})
     if err != nil {
         return nil, err
     }
@@ -55,7 +44,7 @@ func {{.Name}}(
 
 	var retval {{.ResponseName}}
     err = client.MakeRequest(
-        {{if ne $.Config.ContextType "-"}}ctx{{else}}nil{{end}},
+        {{if ne .Config.ContextType "-"}}ctx{{else}}nil{{end}},
         "{{.Name}}",
         `{{.Body}}`,
         &retval,
@@ -63,4 +52,3 @@ func {{.Name}}(
     )
     return &retval, err
 }
-{{end}}
diff --git a/generate/template.go b/generate/template.go
index 984fdf6..6e3ab43 100644
--- a/generate/template.go
+++ b/generate/template.go
@@ -33,8 +33,8 @@ func intRange(n int) []int {
 
 func sub(x, y int) int { return x - y }
 
-// execute executes the given template with the funcs from this generator.
-func (g *generator) execute(tmplRelFilename string, w io.Writer, data interface{}) error {
+// render executes the given template with the funcs from this generator.
+func (g *generator) render(tmplRelFilename string, w io.Writer, data interface{}) error {
 	tmpl := g.templateCache[tmplRelFilename]
 	if tmpl == nil {
 		absFilename := filepath.Join(thisDir, tmplRelFilename)
diff --git a/generate/types.go b/generate/types.go
index 653c0be..e29f32a 100644
--- a/generate/types.go
+++ b/generate/types.go
@@ -202,14 +202,7 @@ func (typ *goStructType) WriteDefinition(w io.Writer, g *generator) error {
 		return nil
 	}
 
-	// TODO(benkraft): Avoid having to enumerate these in advance; just let the
-	// template add them directly.
-	_, err := g.addRef("encoding/json.Unmarshal")
-	if err != nil {
-		return err
-	}
-
-	return g.execute("unmarshal.go.tmpl", w, typ)
+	return g.render("unmarshal.go.tmpl", w, typ)
 }
 
 func (typ *goStructType) Reference() string              { return typ.GoName }
@@ -287,19 +280,7 @@ func (typ *goInterfaceType) WriteDefinition(w io.Writer, g *generator) error {
 
 	// Finally, write the unmarshal-helper, which will be called by struct
 	// fields referencing this type (see goStructType.WriteDefinition).
-	//
-	// TODO(benkraft): Avoid having to enumerate these refs in advance; just
-	// let the template add them directly.
-	_, err := g.addRef("encoding/json.Unmarshal")
-	if err != nil {
-		return err
-	}
-	_, err = g.addRef("fmt.Errorf")
-	if err != nil {
-		return err
-	}
-
-	return g.execute("unmarshal_helper.go.tmpl", w, typ)
+	return g.render("unmarshal_helper.go.tmpl", w, typ)
 }
 
 func (typ *goInterfaceType) Reference() string              { return typ.GoName }


Test patch:
diff --git a/generate/testdata/snapshots/TestGenerateWithConfig-ClientGetterCustomContext-testdata-queries-generated.go b/generate/testdata/snapshots/TestGenerateWithConfig-ClientGetterCustomContext-testdata-queries-generated.go
index 107279c..098bc98 100644
--- a/generate/testdata/snapshots/TestGenerateWithConfig-ClientGetterCustomContext-testdata-queries-generated.go
+++ b/generate/testdata/snapshots/TestGenerateWithConfig-ClientGetterCustomContext-testdata-queries-generated.go
@@ -3,9 +3,14 @@ package queries
 // Code generated by github.com/Khan/genqlient, DO NOT EDIT.
 
 import (
+	"context"
+
 	"github.com/Khan/genqlient/internal/testutil"
 )
 
+// Check that context_type from genqlient.yaml implements context.Context.
+var _ context.Context = (testutil.MyContext)(nil)
+
 // SimpleQueryResponse is returned by SimpleQuery on success.
 type SimpleQueryResponse struct {
 	// user looks up a user by some stuff.
diff --git a/generate/testdata/snapshots/TestGenerateWithConfig-CustomContext-testdata-queries-generated.go b/generate/testdata/snapshots/TestGenerateWithConfig-CustomContext-testdata-queries-generated.go
index ff9aa39..3c7a7ef 100644
--- a/generate/testdata/snapshots/TestGenerateWithConfig-CustomContext-testdata-queries-generated.go
+++ b/generate/testdata/snapshots/TestGenerateWithConfig-CustomContext-testdata-queries-generated.go
@@ -3,10 +3,15 @@ package queries
 // Code generated by github.com/Khan/genqlient, DO NOT EDIT.
 
 import (
+	"context"
+
 	"github.com/Khan/genqlient/graphql"
 	"github.com/Khan/genqlient/internal/testutil"
 )
 
+// Check that context_type from genqlient.yaml implements context.Context.
+var _ context.Context = (testutil.MyContext)(nil)
+
 // SimpleQueryResponse is returned by SimpleQuery on success.
 type SimpleQueryResponse struct {
 	// user looks up a user by some stuff.