Allow to merge an array of stages into a pipeline

Author: GromNaN

src/Builder/Pipeline.php

@@ -8,26 +8,29 @@
 use IteratorAggregate;
 use MongoDB\Builder\Type\StageInterface;
 use MongoDB\Exception\InvalidArgumentException;
-use Traversable;
 
 use function array_is_list;
 use function array_merge;
+use function is_array;
 
 /**
  * An aggregation pipeline consists of one or more stages that process documents.
  *
  * @see https://www.mongodb.com/docs/manual/core/aggregation-pipeline/
  *
  * @psalm-immutable
- * @implements IteratorAggregate<StageInterface>
+ * @implements IteratorAggregate<StageInterface|array<string,mixed>|object>
  */
-class Pipeline implements IteratorAggregate
+final class Pipeline implements IteratorAggregate
 {
-    /** @var StageInterface[] */
     private readonly array $stages;
 
-    /** @no-named-arguments */
-    public function __construct(StageInterface|Pipeline ...$stagesOrPipelines)
+    /**
+     * @param StageInterface|Pipeline|list<StageInterface> ...$stagesOrPipelines
+     *
+     * @no-named-arguments
+     */
+    public function __construct(StageInterface|Pipeline|array ...$stagesOrPipelines)
     {
         if (! array_is_list($stagesOrPipelines)) {
             throw new InvalidArgumentException('Named arguments are not supported for pipelines');
@@ -36,7 +39,9 @@ public function __construct(StageInterface|Pipeline ...$stagesOrPipelines)
         $stages = [];
 
         foreach ($stagesOrPipelines as $stageOrPipeline) {
-            if ($stageOrPipeline instanceof Pipeline) {
+            if (is_array($stageOrPipeline) && array_is_list($stageOrPipeline)) {
+                $stages = array_merge($stages, $stageOrPipeline);
+            } elseif ($stageOrPipeline instanceof Pipeline) {
                 $stages = array_merge($stages, $stageOrPipeline->stages);
             } else {
                 $stages[] = $stageOrPipeline;
@@ -46,8 +51,7 @@ public function __construct(StageInterface|Pipeline ...$stagesOrPipelines)
         $this->stages = $stages;
     }
 
-    /** @return Traversable<StageInterface> */
-    public function getIterator(): Traversable
+    public function getIterator(): ArrayIterator
     {
         return new ArrayIterator($this->stages);
     }

tests/Builder/PipelineTest.php

@@ -26,15 +26,16 @@ public function testMergingPipeline(): void
     {
         $stages = array_map(
             fn (int $i) => $this->createMock(StageInterface::class),
-            range(0, 5),
+            range(0, 7),
         );
 
         $pipeline = new Pipeline(
             $stages[0],
             $stages[1],
             new Pipeline($stages[2], $stages[3]),
-            $stages[4],
-            new Pipeline($stages[5]),
+            [$stages[4], $stages[5]],
+            new Pipeline($stages[6]),
+            [$stages[7]],
         );
 
         $this->assertSame($stages, iterator_to_array($pipeline));