From 6c9c18746b41bf6e8fcba593918c4551ef5c764a Mon Sep 17 00:00:00 2001
From: Wim Pomp <w.pomp@nki.nl>
Date: Mon, 24 May 2021 18:59:32 +0200
Subject: [PATCH] - Make chunks a class, so we can do len(chunks(...))

---
 parfor/__init__.py | 42 ++++++++++++++++++++++++++----------------
 setup.py           |  2 +-
 2 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/parfor/__init__.py b/parfor/__init__.py
index 73a44f7..3047259 100644
--- a/parfor/__init__.py
+++ b/parfor/__init__.py
@@ -116,7 +116,7 @@ def dumps(obj, protocol=None, byref=None, fmode=None, recurse=True, **kwds):
     return file.getvalue()
 
 
-def chunks(*args, **kwargs):
+class chunks():
     """ Yield successive chunks from lists.
         Usage: chunks(s, list0, list1, ...)
                chunks(list0, list1, ..., s=s)
@@ -125,21 +125,31 @@ def chunks(*args, **kwargs):
         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 = 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]
+
+    def __init__(self, *args, **kwargs):
+        N = min(*[len(a) for a in args]) if len(args) > 1 else len(args[0])
+        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 = min(*[len(a) for a in args]) if len(args) > 1 else len(args[0])
+            n = round(N / s)
+        self.args = args
+        self.A = len(args) == 1
+        self.N = N
+        self.n = max(1, min(N, n))
+
+    def __iter__(self):
+        for i in range(self.n):
+            p, q = (i * self.N // self.n), ((i + 1) * self.N // self.n)
+            yield self.args[0][p:q] if self.A else [a[p:q] for a in self.args]
+
+    def __len__(self):
+        return self.n
 
 
 class tqdmm(tqdm):
diff --git a/setup.py b/setup.py
index a8ba0c2..7faa8cf 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
 
 setuptools.setup(
     name="parfor",
-    version="2021.5.0",
+    version="2021.5.1",
     author="Wim Pomp",
     author_email="wimpomp@gmail.com",
     description="A package to mimic the use of parfor as done in Matlab.",