chore: more specific return type for chunk() method (#351)

* chore: more specific return type for chunk() method * chore(sa): fix psalm type infer for Chunk operation --------- Co-authored-by: Yevhen Sidelnyk <yevhen.sidelnyk@stfalcon.com>
loophp · Dec 11, 2024 · 3cfae9d · 3cfae9d
1 parent af76bb7
commit 3cfae9d
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 20 deletions.
diff --git a/src/Contract/Operation/Chunkable.php b/src/Contract/Operation/Chunkable.php
@@ -17,7 +17,7 @@ interface Chunkable
      *
      * @see https://loophp-collection.readthedocs.io/en/stable/pages/api.html#chunk
      *
-     * @return Collection<int, list<T>>
+     * @return Collection<int, non-empty-list<T>>
      */
     public function chunk(int ...$sizes): Collection;
 }
diff --git a/src/Operation/Chunk.php b/src/Operation/Chunk.php
@@ -18,47 +18,45 @@
 final class Chunk extends AbstractOperation
 {
     /**
-     * @return Closure(int...): Closure(iterable<TKey, T>): Generator<int, list<T>>
+     * @return Closure(int...): Closure(iterable<TKey, T>): Generator<int, non-empty-list<T>>
      */
     public function __invoke(): Closure
     {
         return
             /**
-             * @return Closure(iterable<TKey, T>): Generator<int, list<T>>
+             * @return Closure(iterable<TKey, T>): Generator<int, non-empty-list<T>>
              */
             static fn (int ...$sizes): Closure =>
                 /**
                  * @param iterable<TKey, T> $iterable
                  *
-                 * @return Generator<int, list<T>>
+                 * @return Generator<int, non-empty-list<T>>
                  */
                 static function (iterable $iterable) use ($sizes): Generator {
                     $sizesCount = count($sizes);
-                    $i = 0;
-                    $values = [];
+                    $chunkIndex = 0;
+                    $chunk = [];
 
                     foreach ($iterable as $value) {
-                        $size = $sizes[$i % $sizesCount];
+                        $size = $sizes[$chunkIndex % $sizesCount];
 
                         if (0 >= $size) {
                             return;
                         }
 
-                        if (count($values) !== $size) {
-                            $values[] = $value;
+                        $chunk[] = $value;
 
-                            continue;
-                        }
-
-                        ++$i;
+                        if (count($chunk) >= $size) {
+                            ++$chunkIndex;
 
-                        yield $values;
+                            yield $chunk;
 
-                        $values = [$value];
+                            $chunk = [];
+                        }
                     }
 
-                    if ([] !== $values) {
-                        yield $values;
+                    if ([] !== $chunk) {
+                        yield $chunk;
                     }
                 };
     }

diff --git a/tests/static-analysis/chunk.php b/tests/static-analysis/chunk.php
@@ -8,18 +8,18 @@
 use loophp\collection\Contract\Collection as CollectionInterface;
 
 /**
- * @param CollectionInterface<int, list<int>> $collection
+ * @param CollectionInterface<int, non-empty-list<int>> $collection
  */
 function chunk_checkList(CollectionInterface $collection): void {}
 /**
- * @param CollectionInterface<int, list<string>> $collection
+ * @param CollectionInterface<int, non-empty-list<string>> $collection
  */
 function chunk_checkMap(CollectionInterface $collection): void {}
 
 chunk_checkList(Collection::fromIterable(range(0, 6))->chunk(1));
 chunk_checkList(Collection::fromIterable(range(0, 6))->chunk(1, 2));
 
-// These should work but Psalm restricts type to list<'baz'|'taz'>
+// These should work but Psalm restricts type to non-empty-list<'baz'|'taz'>
 /** @psalm-suppress InvalidArgument */
 chunk_checkMap(Collection::fromIterable(['foo' => 'bar', 'baz' => 'taz'])->chunk(1));
 /** @psalm-suppress InvalidArgument */