From 6813c2d8c1d744ee3db22488d619762db989ba0b Mon Sep 17 00:00:00 2001 From: JackyHsu-X Date: Thu, 22 Oct 2026 23:34:41 +0300 Subject: [PATCH] Reject nonpositive chunksizes in pool imap methods Signed-off-by: JackyHsu-X --- CHANGES.txt | 6 ++++++ billiard/pool.py | 10 ++++++++-- t/unit/test_pool.py | 33 +++++++++++++++++++++++++++++++++ 3 files changed, 47 insertions(+), 2 deletions(-) diff --git a/CHANGES.txt b/CHANGES.txt index dd6bf33..7f4a0e6 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -1,3 +1,9 @@ +4.4.0 - unreleased +------------------ +- Reject non-positive chunksize in Pool.imap() and Pool.imap_unordered(), + including when Python runs with optimization enabled. + + 4.3.1 - 2026-10-05 ------------------ - Reject bool for Pool numeric args (processes, timeouts, limits) - #460 diff --git a/billiard/pool.py b/billiard/pool.py index dd77063..6d407cd 100644 --- a/billiard/pool.py +++ b/billiard/pool.py @@ -1522,6 +1522,10 @@ def imap(self, func, iterable, chunksize=1, lost_worker_timeout=None): if self._state != RUN: return lost_worker_timeout = lost_worker_timeout or self.lost_worker_timeout + if chunksize < 1: + raise ValueError( + 'Chunksize must be 1+, not {0:n}'.format(chunksize), + ) if chunksize == 1: result = IMapIterator(self._cache, lost_worker_timeout=lost_worker_timeout) @@ -1532,7 +1536,6 @@ def imap(self, func, iterable, chunksize=1, lost_worker_timeout=None): )) return result else: - assert chunksize > 1 task_batches = Pool._get_tasks(func, iterable, chunksize) result = IMapIterator(self._cache, lost_worker_timeout=lost_worker_timeout) @@ -1551,6 +1554,10 @@ def imap_unordered(self, func, iterable, chunksize=1, if self._state != RUN: return lost_worker_timeout = lost_worker_timeout or self.lost_worker_timeout + if chunksize < 1: + raise ValueError( + 'Chunksize must be 1+, not {0:n}'.format(chunksize), + ) if chunksize == 1: result = IMapUnorderedIterator( self._cache, lost_worker_timeout=lost_worker_timeout, @@ -1562,7 +1569,6 @@ def imap_unordered(self, func, iterable, chunksize=1, )) return result else: - assert chunksize > 1 task_batches = Pool._get_tasks(func, iterable, chunksize) result = IMapUnorderedIterator( self._cache, lost_worker_timeout=lost_worker_timeout, diff --git a/t/unit/test_pool.py b/t/unit/test_pool.py index e6ded21..033c9a2 100644 --- a/t/unit/test_pool.py +++ b/t/unit/test_pool.py @@ -69,6 +69,39 @@ def announce_pid_then_fail_in_cleanup(marker): class test_pool: + @pytest.mark.parametrize('method', ['imap', 'imap_unordered']) + @pytest.mark.parametrize('chunksize', [0, -1]) + @pytest.mark.parametrize('items', [[], [1, 2]]) + def test_imap_rejects_nonpositive_chunksize(self, method, chunksize, items): + pool = Mock(_state=billiard.pool.RUN, _cache={}) + iterable = iter(items) + + with pytest.raises(ValueError, match=r'Chunksize must be 1\+'): + getattr(billiard.pool.Pool, method)( + pool, simple_task, iterable, chunksize=chunksize, + ) + + assert list(iterable) == items + assert not pool._cache + pool._taskqueue.put.assert_not_called() + + @pytest.mark.parametrize('method', ['imap', 'imap_unordered']) + @pytest.mark.parametrize('chunksize', [1, 2]) + def test_imap_valid_chunksize_after_rejection(self, method, chunksize): + pool = Mock(_state=billiard.pool.RUN, _cache={}) + imap = getattr(billiard.pool.Pool, method) + with pytest.raises(ValueError): + imap(pool, simple_task, [1, 2], chunksize=0) + result = imap(pool, simple_task, [1, 2, 3], chunksize) + (tasks, set_length), = pool._taskqueue.put.call_args.args + count = 0 + for _, (job, index, func, args, kwargs) in tasks: + pool._cache[job]._set(index, (True, func(*args, **kwargs))) + count += 1 + set_length(count) + assert sorted(result) == [2, 4, 6] + assert not pool._cache + def test_memory_error_from_callback_propagates(self): def callback(value): raise MemoryError(value)