From 25a2b45ccc2f269f46ddd5e0f599eedaeaf7e28d Mon Sep 17 00:00:00 2001
From: Ben Kraft <benkraft@khanacademy.org>
Date: Fri, 9 Apr 2021 09:53:53 -0700
Subject: [PATCH] redo comment parsing to set up for field-level config

---
 README.md                                     |  2 +-
 generate/comments.go                          | 47 +++++++++++++++++++
 generate/config.go                            |  3 ++
 generate/generate.go                          | 30 +++++-------
 .../testdata/queries/SimpleMutation.graphql   |  2 +
 5 files changed, 64 insertions(+), 20 deletions(-)
 create mode 100644 generate/comments.go

diff --git a/README.md b/README.md
index e3736d0..8694dbf 100644
--- a/README.md
+++ b/README.md
@@ -101,7 +101,7 @@ Generated code customization:
 - add flag(s) to make a field use a pointer (for optionality or perf; see DESIGN)
 - omitempty-like thing for optional arguments
 - collapsing -- should be able to have `mutation { myMutation { error { code } } }` just return `(code string, err error)`
-- map a GraphQL type to a particular Go type (if you want to use a named type for some string, say)
+- map a field to a particular Go type (if you want to use a named type for some string, say)
 - specify a particular name for a Go type, or for the helper function, or whatever else
 - include full query in generated godoc
 
diff --git a/generate/comments.go b/generate/comments.go
new file mode 100644
index 0000000..ff2faeb
--- /dev/null
+++ b/generate/comments.go
@@ -0,0 +1,47 @@
+package generate
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/vektah/gqlparser/v2/ast"
+	"github.com/vektah/gqlparser/v2/parser"
+)
+
+func (g *generator) parsePrecedingComment(pos *ast.Position) (comment string, directives []*ast.Directive, err error) {
+	var commentLines []string
+	sourceLines := strings.Split(pos.Src.Input, "\n")
+	for i := pos.Line - 1; i > 0; i-- {
+		line := strings.TrimSpace(sourceLines[i-1])
+		trimmed := strings.TrimSpace(strings.TrimPrefix(line, "#"))
+		if strings.HasPrefix(line, "# @genqlient") {
+			directive, err := parseDirective(trimmed, pos)
+			if err != nil {
+				return "", nil, err
+			}
+			directives = append(directives, directive)
+		} else if strings.HasPrefix(line, "#") {
+			commentLines = append(commentLines, trimmed)
+		} else {
+			break
+		}
+	}
+
+	reverse(commentLines)
+
+	return strings.TrimSpace(strings.Join(commentLines, "\n")), directives, nil
+}
+
+func parseDirective(line string, pos *ast.Position) (*ast.Directive, error) {
+	// HACK: parse the "directive" by making a fake query containing it.
+	fakeQuery := fmt.Sprintf("query %v { field }", line)
+	doc, err := parser.ParseQuery(&ast.Source{
+		Name: fmt.Sprintf("@genqlient directive at %v:%v:%v",
+			pos.Src.Name, pos.Line, pos.Column),
+		Input: fakeQuery,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return doc.Operations[0].Directives[0], nil
+}
diff --git a/generate/config.go b/generate/config.go
index d1c8aed..0b84f34 100644
--- a/generate/config.go
+++ b/generate/config.go
@@ -76,6 +76,9 @@ type Config struct {
 	// 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.
+	// TODO: figure out if it makes sense to say you can use these for
+	// non-scalar types; technically it should just work, but what if you
+	// didn't request the right fields?
 	Scalars map[string]string `yaml:"scalars"`
 
 	// Set automatically to the filename of the config file itself.
diff --git a/generate/generate.go b/generate/generate.go
index 13e7ae5..8203732 100644
--- a/generate/generate.go
+++ b/generate/generate.go
@@ -120,24 +120,6 @@ func (g *generator) getArgument(opName string, arg *ast.VariableDefinition) (arg
 	}, nil
 }
 
-func (g *generator) getDocComment(op *ast.OperationDefinition) string {
-	var commentLines []string
-	sourceLines := strings.Split(op.Position.Src.Input, "\n")
-	for i := op.Position.Line - 1; i > 0; i-- {
-		line := strings.TrimSpace(sourceLines[i-1])
-		if strings.HasPrefix(line, "#") && !strings.HasPrefix(line, "# @genqlient") {
-			commentLines = append(commentLines,
-				"// "+strings.TrimSpace(strings.TrimPrefix(line, "#")))
-		} else {
-			break
-		}
-	}
-
-	reverse(commentLines)
-
-	return strings.Join(commentLines, "\n")
-}
-
 func (g *generator) addOperation(op *ast.OperationDefinition) error {
 	if op.Name == "" {
 		return fmt.Errorf("operations must have operation-names")
@@ -165,10 +147,20 @@ func (g *generator) addOperation(op *ast.OperationDefinition) error {
 		return err
 	}
 
+	commentLines, _, err := g.parsePrecedingComment(op.Position)
+	if err != nil {
+		return err
+	}
+
+	var docComment string
+	if len(commentLines) > 0 {
+		docComment = "// " + strings.ReplaceAll(commentLines, "\n", "\n// ")
+	}
+
 	g.Operations = append(g.Operations, operation{
 		Type: op.Operation,
 		Name: op.Name,
-		Doc:  g.getDocComment(op),
+		Doc:  docComment,
 		// The newline just makes it format a little nicer.
 		Body:           "\n" + builder.String(),
 		Args:           args,
diff --git a/generate/testdata/queries/SimpleMutation.graphql b/generate/testdata/queries/SimpleMutation.graphql
index d2afc81..a5c3c00 100644
--- a/generate/testdata/queries/SimpleMutation.graphql
+++ b/generate/testdata/queries/SimpleMutation.graphql
@@ -1,3 +1,5 @@
+# @genqlient
+#
 # SimpleMutation creates a user.
 #
 # It has a long doc-comment, to test that we handle that correctly.