Skip to content

Commit 4b33a9d

Browse files
gh-158004: preserve asyncio.as_completed() items after anext cancellation (#158005)
Co-authored-by: Kumar Aditya <kumaraditya@python.org>
1 parent b12fcff commit 4b33a9d

3 files changed

Lines changed: 47 additions & 1 deletion

File tree

‎Lib/asyncio/tasks.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -636,7 +636,11 @@ async def _wait_for_one(self, resolve=False):
636636
# Wait for the next future to be done and return it unless resolve is
637637
# set, in which case return either the result of the future or raise
638638
# an exception.
639-
f = await self._done.get()
639+
try:
640+
f = await self._done.get()
641+
except:
642+
self._todo_left += 1
643+
raise
640644
if f is None:
641645
# Dummy value from _handle_timeout().
642646
raise exceptions.TimeoutError

‎Lib/test/test_asyncio/test_tasks.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1702,6 +1702,46 @@ async def coro(i):
17021702
)
17031703
self.assertCountEqual(results, (0, 1, 2, 3))
17041704

1705+
def test_as_completed_async_iterator_cancelled_anext(self):
1706+
async def main():
1707+
loop = asyncio.get_running_loop()
1708+
a = loop.create_future()
1709+
b = loop.create_future()
1710+
iterator = asyncio.as_completed([a, b])
1711+
1712+
waiter = asyncio.create_task(anext(iterator))
1713+
await asyncio.sleep(0)
1714+
waiter.cancel()
1715+
with self.assertRaises(asyncio.CancelledError):
1716+
await waiter
1717+
1718+
a.set_result('a')
1719+
b.set_result('b')
1720+
return [await f async for f in iterator]
1721+
1722+
results = self.loop.run_until_complete(main())
1723+
self.assertCountEqual(results, ['a', 'b'])
1724+
1725+
def test_as_completed_plain_iterator_cancelled_wait(self):
1726+
async def main():
1727+
loop = asyncio.get_running_loop()
1728+
a = loop.create_future()
1729+
b = loop.create_future()
1730+
iterator = asyncio.as_completed([a, b])
1731+
1732+
waiter = asyncio.create_task(next(iterator))
1733+
await asyncio.sleep(0)
1734+
waiter.cancel()
1735+
with self.assertRaises(asyncio.CancelledError):
1736+
await waiter
1737+
1738+
a.set_result('a')
1739+
b.set_result('b')
1740+
return [await c for c in iterator]
1741+
1742+
results = self.loop.run_until_complete(main())
1743+
self.assertCountEqual(results, ['a', 'b'])
1744+
17051745
def test_as_completed_reverse_wait(self):
17061746
# Tests the plain iterator style of as_completed iteration to
17071747
# ensure that the first future awaited resolves to the first
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix :func:`asyncio.as_completed` to continue yielding every awaitable after
2+
a pending iteration step is cancelled.

0 commit comments

Comments
 (0)