diff --git a/passerelle/utils/__init__.py b/passerelle/utils/__init__.py index c4be935b..3633bdec 100644 --- a/passerelle/utils/__init__.py +++ b/passerelle/utils/__init__.py @@ -355,7 +355,10 @@ def batch(iterable, size): batchiter = islice(sourceiter, size) # call next() at least one time to advance, if the caller does not # consume the returned iterators, sourceiter will never be exhausted. - yield chain([batchiter.next()], batchiter) + try: + yield chain([next(batchiter)], batchiter) + except StopIteration: + return # legacy import, other modules keep importing to_json from passerelle.utils from .jsonresponse import to_json diff --git a/tests/test_utils_batch.py b/tests/test_utils_batch.py new file mode 100644 index 00000000..da146431 --- /dev/null +++ b/tests/test_utils_batch.py @@ -0,0 +1,7 @@ +from passerelle.utils import batch + + +def test_batch(): + data = range(10) + for i, d in enumerate(batch(data, 1)): + assert list(d) == [i]