Feat: iterable_chunk() (#50)

Author: bpolaszek

README.md

@@ -16,6 +16,7 @@ This package provides functions to work with [iterables](https://wiki.php.net/rf
 - [iterable_reduce()](#iterable_reduce)
 - [iterable_filter()](#iterable_filter)
 - [iterable_values()](#iterable_values)
+- [iterable_chunk()](#iterable_chunk)
 
 iterable_to_array()
 -------------------
@@ -170,6 +171,33 @@ foreach (iterable_values($generator()) as $key => $value) {
 }
 ```
 
+iterable_chunk()
+--------------
+
+Here's an `array_chunk`-like function that also works with a `Traversable`.
+
+```php
+use function BenTools\IterableFunctions\iterable_chunk;
+
+$fruits = [
+    'banana',
+    'apple',
+    'strawberry',
+    'raspberry',
+    'pineapple',
+]
+$fruits = (fn () => yield from $fruits)()
+iterable_chunk($fruits, 2);
+
+/*
+  [
+    ['banana', 'apple'],
+    ['strawberry', 'raspberry'],
+    ['pineapple'],
+  ]
+ */
+```
+
 Iterable fluent interface
 =========================
 

src/ChunkIterator.php

@@ -0,0 +1,94 @@
+<?php
+
+declare(strict_types=1);
+
+namespace BenTools\IterableFunctions;
+
+use Iterator;
+use IteratorIterator;
+use Traversable;
+
+/**
+ * @internal
+ *
+ * @template TKey
+ * @template TValue
+ *
+ * @implements Iterator<int, array<TKey, TValue>>
+ */
+final class ChunkIterator implements Iterator
+{
+    /** @var Iterator<TKey, TValue> */
+    private Iterator $iterator;
+
+    private int $chunkSize;
+
+    private bool $preserveKeys;
+
+    private int $chunkIndex = 0;
+
+    /** @var array<TKey, TValue> */
+    private array $buffer = [];
+
+    /**
+     * @param Traversable<TKey, TValue> $iterator
+     */
+    public function __construct(
+        Traversable $iterator,
+        int $chunkSize,
+        bool $preserveKeys = false,
+    ) {
+        $this->iterator = $iterator instanceof Iterator ? $iterator : new IteratorIterator($iterator);
+        $this->chunkSize = $chunkSize;
+        $this->preserveKeys = $preserveKeys;
+    }
+
+    public function current(): mixed
+    {
+        return $this->buffer;
+    }
+
+    public function next(): void
+    {
+        $this->fill();
+        $this->chunkIndex++;
+    }
+
+    public function key(): int
+    {
+        return $this->chunkIndex;
+    }
+
+    public function valid(): bool
+    {
+        if ($this->chunkIndex === 0) {
+            $this->fill();
+        }
+
+        return $this->buffer !== [];
+    }
+
+    public function rewind(): void
+    {
+        $this->iterator->rewind();
+        $this->chunkIndex = 0;
+        $this->buffer = [];
+    }
+
+    private function fill(): void
+    {
+        $this->buffer = [];
+        $i = 0;
+        while ($this->iterator->valid() && $i++ < $this->chunkSize) {
+            $current = $this->iterator->current();
+
+            if ($this->preserveKeys) {
+                $this->buffer[$this->iterator->key()] = $current;
+            } else {
+                $this->buffer[] = $current;
+            }
+
+            $this->iterator->next();
+        }
+    }
+}

src/IterableObject.php

@@ -12,6 +12,7 @@
 use IteratorIterator;
 use Traversable;
 
+use function array_chunk;
 use function array_filter;
 use function array_map;
 use function iterator_to_array;
@@ -118,6 +119,16 @@ public function values(): self
         return new self(new WithoutKeysTraversable($this->iterable));
     }
 
+    /** @return iterable<int, array<TKey, TValue>> */
+    public function chunk(int $size): iterable
+    {
+        if ($this->iterable instanceof Traversable) {
+            return new ChunkIterator($this->iterable, $size, $this->preserveKeys);
+        }
+
+        return array_chunk($this->iterable, $size, $this->preserveKeys);
+    }
+
     /** @return Traversable<TKey, TValue> */
     public function getIterator(): Traversable
     {

src/iterable-functions.php

@@ -136,6 +136,21 @@ function iterable_values(iterable $iterable): iterable
     return iterable($iterable)->values();
 }
 
+/**
+ * Split an iterable into chunks
+ *
+ * @param iterable<TKey, TValue> $iterable
+ *
+ * @return iterable<iterable<TKey, TValue>>
+ *
+ * @template TKey
+ * @template TValue
+ */
+function iterable_chunk(iterable $iterable, int $size, bool $preserveKeys = false): iterable
+{
+    return iterable($iterable, $preserveKeys)->chunk($size);
+}
+
 /**
  * @param iterable<TKey, TValue>|null $iterable
  *
@@ -144,7 +159,7 @@ function iterable_values(iterable $iterable): iterable
  * @template TKey
  * @template TValue
  */
-function iterable(?iterable $iterable): IterableObject
+function iterable(?iterable $iterable, bool $preserveKeys = true): IterableObject
 {
-    return new IterableObject($iterable ?? new EmptyIterator());
+    return new IterableObject($iterable ?? new EmptyIterator(), $preserveKeys);
 }

tests/IterableChunkTest.php

@@ -0,0 +1,55 @@
+<?php
+
+declare(strict_types=1);
+
+namespace BenTools\IterableFunctions\Tests;
+
+use function BenTools\IterableFunctions\iterable_chunk;
+use function expect;
+use function it;
+
+it('chunks an iterable', function (iterable $fruits): void {
+    $chunks = iterable_chunk($fruits, 2);
+    $expectedChunks = [
+        ['banana', 'apple'],
+        ['strawberry', 'raspberry'],
+        ['pineapple'],
+    ];
+    expect([...$chunks])->toEqual($expectedChunks);
+})->with(function () {
+    $fruits = [
+        'banana',
+        'apple',
+        'strawberry',
+        'raspberry',
+        'pineapple',
+    ];
+    yield 'array' => [$fruits];
+    yield 'traversable' => [(fn () => yield from $fruits)()];
+});
+
+it('preserves keys', function (iterable $fruits): void {
+    $chunks = iterable_chunk($fruits, 2, true);
+    $expectedChunks = [
+        [
+            'banana' => 0,
+            'apple' => 1,
+        ],
+        [
+            'strawberry' => 2,
+            'raspberry' => 3,
+        ],
+        ['pineapple' => 4],
+    ];
+    expect([...$chunks])->toEqual($expectedChunks);
+})->with(function () {
+    $fruits = [
+        'banana' => 0,
+        'apple' => 1,
+        'strawberry' => 2,
+        'raspberry' => 3,
+        'pineapple' => 4,
+    ];
+    yield 'array' => [$fruits];
+    yield 'traversable' => [(fn () => yield from $fruits)()];
+});