From 95ae32502857a5269df4e5a92838e182f946e685 Mon Sep 17 00:00:00 2001
From: Wim Pomp <w.pomp@nki.nl>
Date: Mon, 24 May 2021 18:19:14 +0200
Subject: [PATCH] - Better chunks function, option for both number of chunks or
 size of chunks.

---
 parfor/__init__.py | 32 +++++++++++++++++++++++---------
 setup.py           |  2 +-
 2 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/parfor/__init__.py b/parfor/__init__.py
index 57c7a60..73a44f7 100644
--- a/parfor/__init__.py
+++ b/parfor/__init__.py
@@ -116,16 +116,30 @@ def dumps(obj, protocol=None, byref=None, fmode=None, recurse=True, **kwds):
     return file.getvalue()
 
 
-def chunks(n, *args):
-    """ Yield successive n-sized chunks from lists. """
+def chunks(*args, **kwargs):
+    """ Yield successive chunks from lists.
+        Usage: chunks(s, list0, list1, ...)
+               chunks(list0, list1, ..., s=s)
+               chunks(list0, list1, ..., n=n)
+        s: size of chunks, might change to optimize devision between chunks
+        n: number of chunks, coerced to 1 <= n <= len(list0)
+        both s and n are given: use n, unless the chunk size would be bigger than s
+    """
+    N = len(args[-1])
+    if 's' in kwargs and 'n' in kwargs:
+        n = kwargs['n'] if N < kwargs['s'] * kwargs['n'] else round(N / kwargs['s'])
+    elif 's' in kwargs:  # size of chunks
+        n = round(N / kwargs['s'])
+    elif 'n' in kwargs:  # number of chunks
+        n = kwargs['n']
+    else:  # size of chunks in 1st argument
+        s, *args = args
+        n = round(N / s)
     A = len(args) == 1
-    N = len(args[0])
-    n = int(round(N/max(1, round(N/n))))
-    for i in range(0, N, n) if N else []:
-        if A:
-            yield args[0][i:i+n]
-        else:
-            yield [a[i:i+n] for a in args]
+    n = max(1, min(N, n))
+    for i in range(n):
+        p, q = (i * N // n), ((i + 1) * N // n)
+        yield args[0][p:q] if A else [a[p:q] for a in args]
 
 
 class tqdmm(tqdm):
diff --git a/setup.py b/setup.py
index e368fbe..a8ba0c2 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
 
 setuptools.setup(
     name="parfor",
-    version="2021.3.2",
+    version="2021.5.0",
     author="Wim Pomp",
     author_email="wimpomp@gmail.com",
     description="A package to mimic the use of parfor as done in Matlab.",