From 3a8301f6c6d18675df865eea20f92960bd1b4cc1 Mon Sep 17 00:00:00 2001
From: Vladimir Razuvaev <vladimir.razuvaev@gmail.com>
Date: Fri, 28 Jul 2017 17:53:57 +0700
Subject: [PATCH] Allow `types` schema option to be callable

---
 src/Config.php | 27 +++++----------------------
 src/Schema.php | 42 +++++++++++++++++++++++++++++++++---------
 2 files changed, 38 insertions(+), 31 deletions(-)

diff --git a/src/Config.php b/src/Config.php
index 9d91b4f..58fe56f 100644
--- a/src/Config.php
+++ b/src/Config.php
@@ -31,7 +31,7 @@ class Config
     public $subscription;
 
     /**
-     * @var Type[]
+     * @var Type[]|callable
      */
     public $types;
 
@@ -88,8 +88,8 @@ class Config
 
             if (isset($options['types'])) {
                 Utils::invariant(
-                    is_array($options['types']),
-                    'Schema types must be array if provided but got: %s',
+                    is_array($options['types']) || is_callable($options['types']),
+                    'Schema types must be array or callable if provided but got: %s',
                     Utils::getVariableType($options['types'])
                 );
                 $config->setTypes($options['types']);
@@ -110,6 +110,7 @@ class Config
                     'Schema type loader must be callable if provided but got: %s',
                     Utils::getVariableType($options['typeLoader'])
                 );
+                $config->setTypeLoader($options['typeLoader']);
             }
 
             if (isset($options['descriptor'])) {
@@ -188,20 +189,11 @@ class Config
     }
 
     /**
-     * @param Type[] $types
+     * @param Type[]|callable $types
      * @return Config
      */
     public function setTypes($types)
     {
-        foreach ($types as $index => $type) {
-            Utils::invariant(
-                $type instanceof Type,
-                'Schema types must be GraphQL\Type\Definition\Type[], but entry at index "%s" is "%s"',
-                $index,
-                Utils::getVariableType($type)
-            );
-        }
-
         $this->types = $types;
         return $this;
     }
@@ -220,15 +212,6 @@ class Config
      */
     public function setDirectives(array $directives)
     {
-        foreach ($directives as $index => $directive) {
-            Utils::invariant(
-                $directive instanceof Directive,
-                'Schema directives must be GraphQL\Type\Definition\Directive[] if provided but but entry at index "%s" is "%s"',
-                $index,
-                Utils::getVariableType($directive)
-            );
-        }
-
         $this->directives = $directives;
         return $this;
     }
diff --git a/src/Schema.php b/src/Schema.php
index 51c34bb..9ef4af8 100644
--- a/src/Schema.php
+++ b/src/Schema.php
@@ -224,19 +224,43 @@ class Schema
 
     private function collectAllTypes()
     {
-        $initialTypes = [
-            $this->config->query,
-            $this->config->mutation,
-            $this->config->subscription,
-            Introspection::_schema()
-        ];
-        if (!empty($this->config->types)) {
-            $initialTypes = array_merge($initialTypes, $this->config->types);
-        }
+        $initialTypes = array_merge(
+            [
+                $this->config->query,
+                $this->config->mutation,
+                $this->config->subscription,
+                Introspection::_schema()
+            ],
+            array_values($this->resolvedTypes)
+        );
+
         $typeMap = [];
         foreach ($initialTypes as $type) {
             $typeMap = TypeInfo::extractTypes($type, $typeMap);
         }
+
+        $types = $this->config->types;
+        if (is_callable($types)) {
+            $types = $types();
+
+            Utils::invariant(
+                is_array($types) || $types instanceof \Traversable,
+                'Schema types callable must return array or instance of Traversable but got: %s',
+                Utils::getVariableType($types)
+            );
+        }
+
+        if (!empty($types)) {
+            foreach ($types as $type) {
+                Utils::invariant(
+                    $type instanceof Type,
+                    'Each entry of schema types must be instance of GraphQL\Type\Definition\Type but got: %s',
+                    Utils::getVariableType($types)
+                );
+                $typeMap = TypeInfo::extractTypes($type, $typeMap);
+            }
+        }
+
         return $typeMap + Type::getInternalTypes();
     }