From da89e5eb2515614658044bc0e98be63790bdf188 Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Tue, 28 Jun 2022 18:37:17 +0200
Subject: [PATCH] Add chunk

---
 src/great_ai/utilities/__init__.py |  1 +
 src/great_ai/utilities/chunk.py    | 17 ++++++++++++++++
 tests/utilities/test_chunk.py      | 32 ++++++++++++++++++++++++++++++
 3 files changed, 50 insertions(+)
 create mode 100644 src/great_ai/utilities/chunk.py
 create mode 100644 tests/utilities/test_chunk.py

diff --git a/src/great_ai/utilities/__init__.py b/src/great_ai/utilities/__init__.py
index c06e891..1b119e2 100644
--- a/src/great_ai/utilities/__init__.py
+++ b/src/great_ai/utilities/__init__.py
@@ -1,3 +1,4 @@
+from .chunk import chunk
 from .clean import clean
 from .config_file import ConfigFile, ParseError
 from .evaluate_ranking import evaluate_ranking
diff --git a/src/great_ai/utilities/chunk.py b/src/great_ai/utilities/chunk.py
new file mode 100644
index 0000000..05ac663
--- /dev/null
+++ b/src/great_ai/utilities/chunk.py
@@ -0,0 +1,17 @@
+from typing import Iterable, List, TypeVar
+
+T = TypeVar("T")
+
+
+def chunk(values: Iterable[T], chunk_length: int) -> Iterable[T]:
+    assert chunk_length >= 1
+
+    result: List[T] = []
+    for v in values:
+        result.append(v)
+        if len(result) == chunk_length:
+            yield result
+            result = []
+
+    if len(result) > 0:
+        yield result
diff --git a/tests/utilities/test_chunk.py b/tests/utilities/test_chunk.py
new file mode 100644
index 0000000..2a80579
--- /dev/null
+++ b/tests/utilities/test_chunk.py
@@ -0,0 +1,32 @@
+import unittest
+
+import pytest
+
+from src.great_ai.utilities import chunk
+
+
+class TestChunk(unittest.TestCase):
+    def test_simple(self) -> None:
+        i = [1, 2, 3, 4]
+
+        assert list(chunk(i, 1)) == [[1], [2], [3], [4]]
+        assert list(chunk(i, 2)) == [[1, 2], [3, 4]]
+        assert list(chunk(i, 3)) == [[1, 2, 3], [4]]
+        assert list(chunk(i, 4)) == [[1, 2, 3, 4]]
+        assert list(chunk(i, 5)) == [[1, 2, 3, 4]]
+        assert list(chunk(i, 125)) == [[1, 2, 3, 4]]
+
+    def test_bad_argument(self) -> None:
+        with pytest.raises(AssertionError):
+            list(chunk([], -10))
+
+        with pytest.raises(AssertionError):
+            list(chunk([], 0))
+
+    def test_generator(self) -> None:
+        def my_generator():
+            for i in range(1, 11):
+                yield i
+
+        assert list(chunk(range(1, 11), 3)) == [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10]]
+        assert list(chunk(my_generator(), 3)) == [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10]]