From 5debbc6acfc7aa0b5ace5c96d775401aef4ad85f Mon Sep 17 00:00:00 2001
From: vvakame <vvakame+dev@gmail.com>
Date: Mon, 2 Jul 2018 19:05:27 +0900
Subject: [PATCH] Implement types.yaml parsing

---
 Gopkg.lock             |  8 +++++++-
 Gopkg.toml             |  4 ++++
 codegen/codegen.go     |  6 +++---
 main.go                | 35 +++++++++++++++++++++++++++++++++++
 test/resolvers_test.go |  2 +-
 test/types.json        | 11 -----------
 test/types.yaml        |  8 ++++++++
 7 files changed, 58 insertions(+), 16 deletions(-)
 delete mode 100644 test/types.json
 create mode 100644 test/types.yaml

diff --git a/Gopkg.lock b/Gopkg.lock
index 0480a7363e8..ac07db6dfc1 100644
--- a/Gopkg.lock
+++ b/Gopkg.lock
@@ -127,6 +127,12 @@
   ]
   revision = "ce871d178848e3eea1e8795e5cfb74053dde4bb9"
 
+[[projects]]
+  name = "gopkg.in/yaml.v2"
+  packages = ["."]
+  revision = "5420a8b6744d3b0345ab293f6fcba19c978f1183"
+  version = "v2.2.1"
+
 [[projects]]
   branch = "master"
   name = "sourcegraph.com/sourcegraph/appdash"
@@ -150,6 +156,6 @@
 [solve-meta]
   analyzer-name = "dep"
   analyzer-version = 1
-  inputs-digest = "5275b6a5c88fd233faf5afa8f59f3a36ae9fe6538ee6331fac3c1078fecd6df5"
+  inputs-digest = "ee97741e2ec373079694c0adc517dd7bb0dc3271ec459281fa0209d066f04995"
   solver-name = "gps-cdcl"
   solver-version = 1
diff --git a/Gopkg.toml b/Gopkg.toml
index 69486f6fd11..b12c5375ffe 100644
--- a/Gopkg.toml
+++ b/Gopkg.toml
@@ -15,3 +15,7 @@ required = ["github.com/vektah/dataloaden"]
 [prune]
   go-tests = true
   unused-packages = true
+
+[[constraint]]
+  name = "gopkg.in/yaml.v2"
+  version = "2.2.1"
diff --git a/codegen/codegen.go b/codegen/codegen.go
index 746759e7e6f..25989b8d3f7 100644
--- a/codegen/codegen.go
+++ b/codegen/codegen.go
@@ -51,13 +51,13 @@ func (tm TypeMap) Get(typeName string) *TypeMapEntry {
 }
 
 type TypeMapEntry struct {
-	TypeName   string
-	EntityPath string
+	TypeName   string `yaml:"typeName"`
+	EntityPath string `yaml:"entityPath"`
 	Fields     []TypeMapField
 }
 
 type TypeMapField struct {
-	FieldName string
+	FieldName string `yaml:"fieldName"`
 }
 
 func Generate(cfg Config) error {
diff --git a/main.go b/main.go
index 1e8180269c9..44cd3a08887 100644
--- a/main.go
+++ b/main.go
@@ -6,8 +6,10 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"strings"
 
 	"github.com/vektah/gqlgen/codegen"
+	"gopkg.in/yaml.v2"
 )
 
 var output = flag.String("out", "generated.go", "the file to write to")
@@ -58,6 +60,19 @@ func loadTypeMap() codegen.TypeMap {
 	if *typemap == "" {
 		return nil
 	}
+
+	if strings.HasSuffix(*typemap, ".json") {
+		return loadTypeMapJSON()
+	}
+
+	return loadTypeMapYAML()
+}
+
+func loadTypeMapJSON() codegen.TypeMap {
+	if *typemap == "" {
+		return nil
+	}
+
 	b, err := ioutil.ReadFile(*typemap)
 	if err != nil {
 		fmt.Fprintln(os.Stderr, "unable to open typemap: "+err.Error())
@@ -89,3 +104,23 @@ func loadTypeMap() codegen.TypeMap {
 
 	return typeMap
 }
+
+func loadTypeMapYAML() codegen.TypeMap {
+	if *typemap == "" {
+		return nil
+	}
+
+	b, err := ioutil.ReadFile(*typemap)
+	if err != nil {
+		fmt.Fprintln(os.Stderr, "unable to open typemap: "+err.Error())
+		return nil
+	}
+
+	var typeMap codegen.TypeMap
+	if err = yaml.Unmarshal(b, &typeMap); err != nil {
+		fmt.Fprintln(os.Stderr, "unable to parse typemap: "+err.Error())
+		os.Exit(1)
+	}
+
+	return typeMap
+}
diff --git a/test/resolvers_test.go b/test/resolvers_test.go
index ef3cb9c81bc..220d86f4635 100644
--- a/test/resolvers_test.go
+++ b/test/resolvers_test.go
@@ -1,4 +1,4 @@
-//go:generate gorunpkg github.com/vektah/gqlgen -out generated.go -typemap types.json -models models/generated.go
+//go:generate gorunpkg github.com/vektah/gqlgen -out generated.go -typemap types.yaml -models models/generated.go
 
 package test
 
diff --git a/test/types.json b/test/types.json
deleted file mode 100644
index 0ca4651f136..00000000000
--- a/test/types.json
+++ /dev/null
@@ -1,11 +0,0 @@
-{
-    "Element": {
-        "entityPath": "github.com/vektah/gqlgen/test.Element"
-    },
-    "Viewer": {
-        "entityPath": "github.com/vektah/gqlgen/test.Viewer"
-    },
-    "User": {
-        "entityPath": "remote_api.User"
-    }
-}
diff --git a/test/types.yaml b/test/types.yaml
new file mode 100644
index 00000000000..c0ed0bad498
--- /dev/null
+++ b/test/types.yaml
@@ -0,0 +1,8 @@
+- typeName:   Element
+  entityPath: github.com/vektah/gqlgen/test.Element
+
+- typeName:   Viewer
+  entityPath: github.com/vektah/gqlgen/test.Viewer
+
+- typeName:   User
+  entityPath: remote_api.User