Add test for chunks(), refs #364

sqlite_utils/db.py

@@ -1,4 +1,5 @@
 from .utils import (
+    chunks,
     sqlite3,
     OperationalError,
     suggest_column_types,
@@ -2995,12 +2996,6 @@ class View(Queryable):
         )
 
 
-def chunks(sequence, size):
-    iterator = iter(sequence)
-    for item in iterator:
-        yield itertools.chain([item], itertools.islice(iterator, size - 1))
-
-
 def jsonify_if_needed(value):
     if isinstance(value, decimal.Decimal):
         return float(value)

sqlite_utils/utils.py

@@ -3,6 +3,7 @@ import contextlib
 import csv
 import enum
 import io
+import itertools
 import json
 import os
 from . import recipes
@@ -306,3 +307,9 @@ def _compile_code(code, imports, variable="value"):
         globals[import_.split(".")[0]] = __import__(import_)
     exec(code_o, globals, locals)
     return locals["fn"]
+
+
+def chunks(sequence, size):
+    iterator = iter(sequence)
+    for item in iterator:
+        yield itertools.chain([item], itertools.islice(iterator, size - 1))

tests/test_utils.py

@@ -25,3 +25,18 @@ def test_decode_base64_values(input, expected, should_be_is):
 def test_find_spatialite():
     spatialite = utils.find_spatialite()
     assert spatialite is None or isinstance(spatialite, str)
+
+
+@pytest.mark.parametrize(
+    "size,expected",
+    (
+        (1, [["a"], ["b"], ["c"], ["d"]]),
+        (2, [["a", "b"], ["c", "d"]]),
+        (3, [["a", "b", "c"], ["d"]]),
+        (4, [["a", "b", "c", "d"]]),
+    ),
+)
+def test_chunks(size, expected):
+    input = ["a", "b", "c", "d"]
+    chunks = list(map(list, utils.chunks(input, size)))
+    assert chunks == expected