Merge pull request #322 from renato145/fix_chunked

chunked() getting stuck

Author: jph00

fastcore/basics.py

@@ -210,7 +210,7 @@ def gen(func, seq, cond=true):
 def chunked(it, chunk_sz=None, drop_last=False, n_chunks=None):
     "Return batches from iterator `it` of size `chunk_sz` (or return `n_chunks` total)"
     assert bool(chunk_sz) ^ bool(n_chunks)
-    if n_chunks: chunk_sz = math.ceil(len(it)/n_chunks)
+    if n_chunks: chunk_sz = max(math.ceil(len(it)/n_chunks), 1)
     if not isinstance(it, Iterator): it = iter(it)
     while True:
         res = list(itertools.islice(it, chunk_sz))

nbs/01_basics.ipynb

@@ -1165,7 +1165,7 @@
     "def chunked(it, chunk_sz=None, drop_last=False, n_chunks=None):\n",
     "    \"Return batches from iterator `it` of size `chunk_sz` (or return `n_chunks` total)\"\n",
     "    assert bool(chunk_sz) ^ bool(n_chunks)\n",
-    "    if n_chunks: chunk_sz = math.ceil(len(it)/n_chunks)\n",
+    "    if n_chunks: chunk_sz = max(math.ceil(len(it)/n_chunks), 1)\n",
     "    if not isinstance(it, Iterator): it = iter(it)\n",
     "    while True:\n",
     "        res = list(itertools.islice(it, chunk_sz))\n",
@@ -1197,7 +1197,10 @@
     "\n",
     "t = np.arange(10)\n",
     "test_eq(chunked(t,3),      [[0,1,2], [3,4,5], [6,7,8], [9]])\n",
-    "test_eq(chunked(t,3,True), [[0,1,2], [3,4,5], [6,7,8],    ])"
+    "test_eq(chunked(t,3,True), [[0,1,2], [3,4,5], [6,7,8],    ])\n",
+    "\n",
+    "test_eq(chunked([], 3),          [])\n",
+    "test_eq(chunked([], n_chunks=3), [])"
    ]
   },
   {