From f705f23f971e4d1020f2d132c6f8dece1834f00d Mon Sep 17 00:00:00 2001 From: Pengyu Lee Date: Tue, 20 Oct 2026 20:31:15 +0800 Subject: [PATCH] gh-158004: preserve `asyncio.as_completed()` items after anext cancellation (GH-158005) (cherry picked from commit 4b33a9db8084bb4a5c631db5cb97e4a2e0e33a98) Co-authored-by: Pengyu Lee Co-authored-by: Kumar Aditya --- Lib/asyncio/tasks.py | 6 ++- Lib/test/test_asyncio/test_tasks.py | 40 +++++++++++++++++++ ...-09-23-19-59-03.gh-issue-158004.p68swW.rst | 2 + 3 files changed, 47 insertions(+), 1 deletion(-) create mode 100644 Misc/NEWS.d/next/Library/2026-09-23-19-59-03.gh-issue-158004.p68swW.rst diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 9694ec863a086ad..01228b06fbd082d 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -628,7 +628,11 @@ async def _wait_for_one(self, resolve=False): # Wait for the next future to be done and return it unless resolve is # set, in which case return either the result of the future or raise # an exception. - f = await self._done.get() + try: + f = await self._done.get() + except: + self._todo_left += 1 + raise if f is None: # Dummy value from _handle_timeout(). raise exceptions.TimeoutError diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index 9b64870fae05c4a..f4987ae4a295682 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -1702,6 +1702,46 @@ async def coro(i): ) self.assertCountEqual(results, (0, 1, 2, 3)) + def test_as_completed_async_iterator_cancelled_anext(self): + async def main(): + loop = asyncio.get_running_loop() + a = loop.create_future() + b = loop.create_future() + iterator = asyncio.as_completed([a, b]) + + waiter = asyncio.create_task(anext(iterator)) + await asyncio.sleep(0) + waiter.cancel() + with self.assertRaises(asyncio.CancelledError): + await waiter + + a.set_result('a') + b.set_result('b') + return [await f async for f in iterator] + + results = self.loop.run_until_complete(main()) + self.assertCountEqual(results, ['a', 'b']) + + def test_as_completed_plain_iterator_cancelled_wait(self): + async def main(): + loop = asyncio.get_running_loop() + a = loop.create_future() + b = loop.create_future() + iterator = asyncio.as_completed([a, b]) + + waiter = asyncio.create_task(next(iterator)) + await asyncio.sleep(0) + waiter.cancel() + with self.assertRaises(asyncio.CancelledError): + await waiter + + a.set_result('a') + b.set_result('b') + return [await c for c in iterator] + + results = self.loop.run_until_complete(main()) + self.assertCountEqual(results, ['a', 'b']) + def test_as_completed_reverse_wait(self): # Tests the plain iterator style of as_completed iteration to # ensure that the first future awaited resolves to the first diff --git a/Misc/NEWS.d/next/Library/2026-09-23-19-59-03.gh-issue-158004.p68swW.rst b/Misc/NEWS.d/next/Library/2026-09-23-19-59-03.gh-issue-158004.p68swW.rst new file mode 100644 index 000000000000000..0db4a3d4ca4c4aa --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-09-23-19-59-03.gh-issue-158004.p68swW.rst @@ -0,0 +1,2 @@ +Fix :func:`asyncio.as_completed` to continue yielding every awaitable after +a pending iteration step is cancelled.