From 38976db75ecbf60dad51a5ec4aa9381a2d593aa2 Mon Sep 17 00:00:00 2001 From: JackyHsu-X Date: Fri, 23 Oct 2026 17:38:46 +0300 Subject: [PATCH] Count repeated futures in as_completed --- distributed/client.py | 2 +- distributed/tests/test_as_completed.py | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/distributed/client.py b/distributed/client.py index 5f14cb4911b..5e152716d90 100644 --- a/distributed/client.py +++ b/distributed/client.py @@ -5744,7 +5744,7 @@ def count(self): iterator. """ with self.lock: - return len(self.futures) + len(self.queue.queue) + return sum(self.futures.values()) + len(self.queue.queue) def __repr__(self): return ( diff --git a/distributed/tests/test_as_completed.py b/distributed/tests/test_as_completed.py index b2865a36951..28345dd5ea5 100644 --- a/distributed/tests/test_as_completed.py +++ b/distributed/tests/test_as_completed.py @@ -71,6 +71,29 @@ def test_as_completed_sync(client): assert list(as_completed([])) == [] +@pytest.mark.parametrize("with_results", [False, True]) +@gen_cluster(client=True) +async def test_count_duplicate_futures(c, s, a, b, with_results): + x = c.submit(inc, 1) + y = c.submit(inc, 2) + ac = as_completed([x, x, y], with_results=with_results) + assert ac.count() == 3 + ac.add(x) + assert ac.count() == 4 + ac.update([x]) + assert ac.count() == 5 + + results = [] + for remaining in range(4, -1, -1): + results.append(await ac.__anext__()) + assert ac.count() == remaining + expected_x = (x, 2) if with_results else x + expected_y = (y, 3) if with_results else y + assert results.count(expected_x) == 4 + assert results.count(expected_y) == 1 + assert ac.is_empty() + + def test_as_completed_with_non_futures(client): with pytest.raises(TypeError): list(as_completed([1, 2, 3]))