Preserve keys option on Collection::chunk

cakephp · Sep 29, 2016 · a4cbd70 · a4cbd70
1 parent 3a47ac5
commit a4cbd70
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 5 deletions.
diff --git a/src/Collection/CollectionInterface.php b/src/Collection/CollectionInterface.php
@@ -915,9 +915,10 @@ public function zipWith($items, $callable);
      * ```
      *
      * @param int $chunkSize The maximum size for each chunk
+     * @param boolean $preserveKeys If the keys of the array should be preserved
      * @return \Cake\Collection\CollectionInterface
      */
-    public function chunk($chunkSize);
+    public function chunk($chunkSize, $preserveKeys = false);
 
     /**
      * Returns whether or not there are elements in this collection

diff --git a/src/Collection/CollectionTrait.php b/src/Collection/CollectionTrait.php
@@ -627,16 +627,24 @@ public function zipWith($items, $callable)
      * {@inheritDoc}
      *
      */
-    public function chunk($chunkSize)
+    public function chunk($chunkSize, $preserveKeys = false)
     {
-        return $this->map(function ($v, $k, $iterator) use ($chunkSize) {
-            $values = [$v];
+        return $this->map(function ($v, $k, $iterator) use ($chunkSize, $preserveKeys) {
+            $key = 0;
+            if ($preserveKeys) {
+                $key = $k;
+            }
+            $values = [$key => $v];
             for ($i = 1; $i < $chunkSize; $i++) {
                 $iterator->next();
                 if (!$iterator->valid()) {
                     break;
                 }
-                $values[] = $iterator->current();
+                if ($preserveKeys) {
+                    $values[$iterator->key()] = $iterator->current();
+                } else {
+                    $values[] = $iterator->current();
+                }
             }
 
             return $values;

diff --git a/tests/TestCase/Collection/CollectionTest.php b/tests/TestCase/Collection/CollectionTest.php
@@ -1770,6 +1770,19 @@ public function testChunkNested()
         $this->assertEquals($expected, $chunked);
     }
 
+    /**
+     * Tests the chunk method with preserved keys
+     *
+     * @return void
+     */
+    public function testChunkPreserveKeys()
+    {
+        $collection = new Collection(['a' => 1, 'b' => 2, 'c' => 3, 'd' => 4, 'e' => 5, 'f' => 6, 'g' => 7]);
+        $chunked = $collection->chunk(2, true)->toList();
+        $expected = [['a' => 1, 'b' => 2], ['c' => 3, 'd' => 4], ['e' => 5, 'f' => 6], ['g' => 7]];
+        $this->assertEquals($expected, $chunked);
+    }
+
     /**
      * Tests cartesianProduct
      *
@@ -1961,4 +1974,17 @@ public function testTransposeUnEvenLengthShouldThrowException()
 
         $collection->transpose();
     }
+
+    /**
+     * Tests the chunk method with preserved keys
+     *
+     * @return void
+     */
+    public function testChunkPreserveKeys()
+    {
+        $collection = new Collection(['a' => 1, 'b' => 2, 'c' => 3, 'd' => 4, 'e' => 5, 'f' => 6, 'g' => 7]);
+        $chunked = $collection->chunk(2, true)->toList();
+        $expected = [['a' => 1, 'b' => 2], ['c' => 3, 'd' => 4], ['e' => 5, 'f' => 6], ['g' => 7]];
+        $this->assertEquals($expected, $chunked);
+    }
 }