getTypeMap() does not include input types used in directives

src/type/__tests__/schema-test.js

@@ -10,10 +10,13 @@ import {
   GraphQLInterfaceType,
   GraphQLObjectType,
   GraphQLString,
+  GraphQLInputObjectType,
+  GraphQLDirective,
 } from '../';
 
 import { describe, it } from 'mocha';
 import { expect } from 'chai';
+import { GraphQLList } from '../wrappers';
 
 const InterfaceType = new GraphQLInterfaceType({
   name: 'Interface',
@@ -26,6 +29,37 @@ const ImplementingType = new GraphQLObjectType({
   fields: { fieldName: { type: GraphQLString, resolve: () => '' } },
 });
 
+const DirectiveInputType = new GraphQLInputObjectType({
+  name: 'DirInput',
+  fields: {
+    field: {
+      type: GraphQLString,
+    },
+  },
+});
+
+const WrappedDirectiveInputType = new GraphQLInputObjectType({
+  name: 'WrappedDirInput',
+  fields: {
+    field: {
+      type: GraphQLString,
+    },
+  },
+});
+
+const Directive = new GraphQLDirective({
+  name: 'dir',
+  locations: ['OBJECT'],
+  args: {
+    arg: {
+      type: DirectiveInputType,
+    },
+    argList: {
+      type: new GraphQLList(WrappedDirectiveInputType),
+    },
+  },
+});
+
 const Schema = new GraphQLSchema({
   query: new GraphQLObjectType({
     name: 'Query',
@@ -38,6 +72,7 @@ const Schema = new GraphQLSchema({
       },
     },
   }),
+  directives: [Directive],
 });
 
 describe('Type System: Schema', () => {
@@ -53,4 +88,11 @@ 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');
+    });
+  });
 });

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);
+}