diff --git a/src/ChunkIterator.php b/src/ChunkIterator.php new file mode 100644 index 0000000..7b8d285 --- /dev/null +++ b/src/ChunkIterator.php @@ -0,0 +1,94 @@ +> + */ +final class ChunkIterator implements Iterator +{ + /** @var Iterator */ + private Iterator $iterator; + + private int $chunkSize; + + private bool $preserveKeys; + + private int $chunkIndex = 0; + + /** @var array */ + private array $buffer = []; + + /** + * @param Traversable $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(); + } + } +} diff --git a/src/IterableObject.php b/src/IterableObject.php index c36f7f6..8bc0e74 100644 --- a/src/IterableObject.php +++ b/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> */ + 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 */ public function getIterator(): Traversable { diff --git a/src/iterable-functions.php b/src/iterable-functions.php index 239eb3e..3044042 100644 --- a/src/iterable-functions.php +++ b/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 $iterable + * + * @return iterable> + * + * @template TKey + * @template TValue + */ +function iterable_chunk(iterable $iterable, int $size, bool $preserveKeys = false): iterable +{ + return iterable($iterable, $preserveKeys)->chunk($size); +} + /** * @param iterable|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); } diff --git a/tests/IterableChunkTest.php b/tests/IterableChunkTest.php new file mode 100644 index 0000000..060d2ff --- /dev/null +++ b/tests/IterableChunkTest.php @@ -0,0 +1,55 @@ +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)()]; +});