diff --git a/src/DataLoader.php b/src/DataLoader.php index c66ae22..cc6d911 100644 --- a/src/DataLoader.php +++ b/src/DataLoader.php @@ -116,6 +116,9 @@ public function loadMany($keys) if (!is_array($keys) && !$keys instanceof \Traversable) { throw new \InvalidArgumentException(sprintf('The "%s" method must be called with Array but got: %s.', __METHOD__, gettype($keys))); } + if ($keys instanceof \Traversable) { + $keys = iterator_to_array($keys, false); + } return $this->getPromiseAdapter()->createAll(array_map( function ($key) { return $this->load($key); diff --git a/tests/DataLoadTestCase.php b/tests/DataLoadTestCase.php index 74849f0..25e0077 100644 --- a/tests/DataLoadTestCase.php +++ b/tests/DataLoadTestCase.php @@ -51,6 +51,37 @@ public function testSupportsLoadingMultipleKeysInOneCall() $this->assertEquals([], DataLoader::await($promiseEmpty)); } + /** + * @group primary-api + */ + public function testSupportsLoadingMultipleKeysFromAnArrayIterator() + { + list($identityLoader, $loadCalls) = self::idLoader(); + + $promiseAll = $identityLoader->loadMany(new \ArrayIterator([2 => 'A', 0 => 'B', 1 => 'C'])); + + $this->assertEquals(['A', 'B', 'C'], DataLoader::await($promiseAll)); + $this->assertEquals([['A', 'B', 'C']], $loadCalls->getArrayCopy()); + } + + /** + * @group primary-api + */ + public function testSupportsLoadingEveryKeyFromAGenerator() + { + list($identityLoader, $loadCalls) = self::idLoader(); + $keys = (function () { + yield 1 => 'A'; + yield 1 => 'B'; + yield 0 => 'C'; + })(); + + $promiseAll = $identityLoader->loadMany($keys); + + $this->assertEquals(['A', 'B', 'C'], DataLoader::await($promiseAll)); + $this->assertEquals([['A', 'B', 'C']], $loadCalls->getArrayCopy()); + } + /** * @group primary-api */