From afe4735694ab9ac4e306f20f5485d4496827e1e4 Mon Sep 17 00:00:00 2001
From: Jan Melcher <jan.melcher@aeb.com>
Date: Wed, 3 Jan 2018 09:14:14 +0100
Subject: [PATCH] Include directive argument types in getTypeMap()

---
 src/type/__tests__/schema-test.js | 14 ++++++++++++++
 src/type/schema.js                | 19 ++++++++++++++++++-
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/type/__tests__/schema-test.js b/src/type/__tests__/schema-test.js
index f76ce396bc..ab48895d12 100644
--- a/src/type/__tests__/schema-test.js
+++ b/src/type/__tests__/schema-test.js
@@ -16,6 +16,7 @@ import {
 
 import { describe, it } from 'mocha';
 import { expect } from 'chai';
+import { GraphQLList } from '../wrappers';
 
 const InterfaceType = new GraphQLInterfaceType({
   name: 'Interface',
@@ -37,6 +38,15 @@ const DirectiveInputType = new GraphQLInputObjectType({
   },
 });
 
+const WrappedDirectiveInputType = new GraphQLInputObjectType({
+  name: 'WrappedDirInput',
+  fields: {
+    field: {
+      type: GraphQLString,
+    },
+  },
+});
+
 const Directive = new GraphQLDirective({
   name: 'dir',
   locations: ['OBJECT'],
@@ -44,6 +54,9 @@ const Directive = new GraphQLDirective({
     arg: {
       type: DirectiveInputType,
     },
+    argList: {
+      type: new GraphQLList(WrappedDirectiveInputType),
+    },
   },
 });
 
@@ -79,6 +92,7 @@ describe('Type System: Schema', () => {
   describe('Type Map', () => {
     it('includes input types only used in directives', () => {
       expect(Schema.getTypeMap()).to.include.key('DirInput');
+      expect(Schema.getTypeMap()).to.include.key('WrappedDirInput');
     });
   });
 });
diff --git a/src/type/schema.js b/src/type/schema.js
index d5526443a0..67a4a7d254 100644
--- a/src/type/schema.js
+++ b/src/type/schema.js
@@ -13,6 +13,7 @@ import {
   isUnionType,
   isInputObjectType,
   isWrappingType,
+  getNamedType,
 } from './definition';
 import type {
   GraphQLType,
@@ -22,7 +23,11 @@ import type {
   GraphQLInterfaceType,
 } from './definition';
 import type { SchemaDefinitionNode } from '../language/ast';
-import { GraphQLDirective, specifiedDirectives } from './directives';
+import {
+  GraphQLDirective,
+  isDirective,
+  specifiedDirectives,
+} from './directives';
 import type { GraphQLError } from '../error/GraphQLError';
 import { __Schema } from './introspection';
 import find from '../jsutils/find';
@@ -121,6 +126,10 @@ export class GraphQLSchema {
       initialTypes = initialTypes.concat(types);
     }
 
+    initialTypes = initialTypes.concat(
+      ...this._directives.map(directive => getDirectiveArgTypes(directive)),
+    );
+
     this._typeMap = initialTypes.reduce(
       typeMapReducer,
       (Object.create(null): TypeMap),
@@ -269,3 +278,11 @@ function typeMapReducer(map: TypeMap, type: ?GraphQLType): TypeMap {
 
   return reducedMap;
 }
+
+function getDirectiveArgTypes(directive: GraphQLDirective) {
+  // directives are not validated until validateSchema() is called, so be defensive
+  if (!isDirective(directive)) {
+    return [];
+  }
+  return directive.args.map(arg => arg.type);
+}