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

@@ -22,7 +22,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,10 +125,17 @@ export class GraphQLSchema {
       initialTypes = initialTypes.concat(types);
     }
 
-    this._typeMap = initialTypes.reduce(
-      typeMapReducer,
-      (Object.create(null): TypeMap),
-    );
+    // Keep track of all types referenced within the schema.
+    let typeMap: TypeMap = Object.create(null);
+
+    // First by deeply visiting all initial types.
+    typeMap = initialTypes.reduce(typeMapReducer, typeMap);
+
+    // Then by deeply visiting all directive types.
+    typeMap = this._directives.reduce(typeMapDirectiveReducer, typeMap);
+
+    // Storing the resulting map for reference by the schema.
+    this._typeMap = typeMap;
 
     // Keep track of all implementations by interface name.
     this._implementations = Object.create(null);
@@ -269,3 +280,17 @@ function typeMapReducer(map: TypeMap, type: ?GraphQLType): TypeMap {
 
   return reducedMap;
 }
+
+function typeMapDirectiveReducer(
+  map: TypeMap,
+  directive: ?GraphQLDirective,
+): TypeMap {
+  // Directives are not validated until validateSchema() is called.
+  if (!isDirective(directive)) {
+    return map;
+  }
+  return directive.args.reduce(
+    (_map, arg) => typeMapReducer(_map, arg.type),
+    map,
+  );
+}