fix: array callables + collect()

src/functions.php

@@ -23,7 +23,7 @@
  * @param null|iterable<K,V>|callable():iterable<K,V> $source
  *
  * @return Collection<V,K>
- * @phpstan-return ($source is null ? Collection<never,never> : ($source is array ? ArrayCollection<V> : ($source is DoctrineCollection<K&array-key,V> ? DoctrineBridgeCollection<V> : Collection<V,K>)))
+ * @phpstan-return ($source is null ? Collection<never,never> : ($source is callable ? Collection<V,K> : ($source is array ? ArrayCollection<V> : ($source is DoctrineCollection<K&array-key,V> ? DoctrineBridgeCollection<V> : Collection<V,K>))))
  */
 function collect(iterable|callable|null $source = null): Collection
 {
@@ -35,6 +35,10 @@ function collect(iterable|callable|null $source = null): Collection
         return new DoctrineBridgeCollection($source);
     }
 
+    if (\is_callable($source)) {
+        return new LazyCollection($source(...));
+    }
+
     if (\is_array($source)) {
         return new ArrayCollection($source);
     }

stubs/Types.php

@@ -36,6 +36,9 @@ class User
 
 assertType('Zenstruck\Collection<never, never>', collect());
 assertType('Zenstruck\Collection<never, never>', collect());
+assertType('Zenstruck\Collection\ArrayCollection<User, (int|string)>', collect([new User()]));
+assertType('Zenstruck\Collection<User, int>', collect(fn() => [new User()]));
+assertType('Zenstruck\Collection<User, int>', collect(new \ArrayIterator([new User()])));
 
 /**
  * @param User[]|null $users

tests/FunctionsTest.php

@@ -11,7 +11,11 @@
 
 namespace Zenstruck\Collection\Tests;
 
+use Doctrine\Common\Collections\ArrayCollection as DoctrineArrayCollection;
 use PHPUnit\Framework\TestCase;
+use Zenstruck\Collection\ArrayCollection;
+use Zenstruck\Collection\Doctrine\DoctrineBridgeCollection;
+use Zenstruck\Collection\LazyCollection;
 
 use function Zenstruck\collect;
 
@@ -29,4 +33,111 @@ public function collect(): void
         $this->assertSame([], collect()->eager()->all());
         $this->assertSame([], collect(null)->eager()->all());
     }
+
+    /**
+     * @test
+     */
+    public function collect_arrays_as_array_collection(): void
+    {
+        $this->assertInstanceOf(ArrayCollection::class, collect([]));
+        $this->assertInstanceOf(ArrayCollection::class, collect(['foo' => 'bar']));
+    }
+
+    /**
+     * @test
+     */
+    public function collect_invokes_array_callables(): void
+    {
+        $this->assertInstanceOf(LazyCollection::class, collect([$this, 'items']));
+        $this->assertSame(['foo' => 'bar'], collect([$this, 'items'])->eager()->all());
+
+        $this->assertInstanceOf(LazyCollection::class, collect([self::class, 'staticItems']));
+        $this->assertSame(['foo' => 'bar'], collect([self::class, 'staticItems'])->eager()->all());
+    }
+
+    /**
+     * @test
+     */
+    public function collect_treats_non_callable_arrays_as_arrays(): void
+    {
+        $source = ['foo', 'bar'];
+
+        $this->assertInstanceOf(ArrayCollection::class, collect($source));
+        $this->assertSame($source, collect($source)->eager()->all());
+    }
+
+    /**
+     * @test
+     */
+    public function collect_non_array_callables_as_lazy_collection(): void
+    {
+        $invokable = new class {
+            public function __invoke(): iterable
+            {
+                return ['foo' => 'bar'];
+            }
+        };
+
+        $this->assertInstanceOf(LazyCollection::class, collect(fn() => ['foo' => 'bar']));
+        $this->assertSame(['foo' => 'bar'], collect(fn() => ['foo' => 'bar'])->eager()->all());
+
+        $this->assertInstanceOf(LazyCollection::class, collect($invokable));
+        $this->assertSame(['foo' => 'bar'], collect($invokable)->eager()->all());
+
+        $this->assertInstanceOf(LazyCollection::class, collect('Zenstruck\Collection\Tests\items'));
+        $this->assertSame(['foo' => 'bar'], collect('Zenstruck\Collection\Tests\items')->eager()->all());
+    }
+
+    /**
+     * @test
+     */
+    public function collect_traversables_as_lazy_collection(): void
+    {
+        $collection = collect(new \ArrayIterator(['foo' => 'bar']));
+
+        $this->assertInstanceOf(LazyCollection::class, $collection);
+        $this->assertSame(['foo' => 'bar'], $collection->eager()->all());
+    }
+
+    /**
+     * @test
+     */
+    public function collect_doctrine_collections_as_doctrine_bridge_collection(): void
+    {
+        $collection = collect(new DoctrineArrayCollection(['foo' => 'bar']));
+
+        $this->assertInstanceOf(DoctrineBridgeCollection::class, $collection);
+        $this->assertSame(['foo' => 'bar'], $collection->eager()->all());
+    }
+
+    /**
+     * @test
+     */
+    public function collect_returns_collections_as_is(): void
+    {
+        $collection = ArrayCollection::for(['foo' => 'bar']);
+
+        $this->assertSame($collection, collect($collection));
+    }
+
+    /**
+     * @return array<string,string>
+     */
+    public function items(): array
+    {
+        return ['foo' => 'bar'];
+    }
+
+    /**
+     * @return array<string,string>
+     */
+    public static function staticItems(): array
+    {
+        return ['foo' => 'bar'];
+    }
+}
+
+function items(): iterable
+{
+    return ['foo' => 'bar'];
 }