From 1036ccb55de4abc70837cb46a72ddbb370b8fc94 Mon Sep 17 00:00:00 2001 From: "Miss Islington (bot)" <31488909+miss-islington@users.noreply.github.com> Date: Wed, 26 Aug 2020 10:14:59 -0700 Subject: bpo-32751: Wait for task cancel in asyncio.wait_for() when timeout <= 0 (GH-21895) (GH-21963) When I was fixing bpo-32751 back in GH-7216 I missed the case when *timeout* is zero or negative. This takes care of that. Props to @aaliddell for noticing the inconsistency. (cherry picked from commit c517fc712105c8e5930cb42baaebdbe37fc3e15f) Co-authored-by: Elvis Pranskevichus --- Lib/asyncio/tasks.py | 9 +++++++-- Lib/test/test_asyncio/test_tasks.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 2 deletions(-) (limited to 'Lib') diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index c37f0e1387..7ecec96384 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -445,8 +445,13 @@ async def wait_for(fut, timeout, *, loop=None): if fut.done(): return fut.result() - fut.cancel() - raise exceptions.TimeoutError() + await _cancel_and_wait(fut, loop=loop) + try: + fut.result() + except exceptions.CancelledError as exc: + raise exceptions.TimeoutError() from exc + else: + raise exceptions.TimeoutError() waiter = loop.create_future() timeout_handle = loop.call_later(timeout, _release_waiter, waiter) diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index f9db066ce8..511961c320 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -1131,6 +1131,9 @@ class BaseTaskTests: nonlocal task_done try: await asyncio.sleep(0.2) + except asyncio.CancelledError: + await asyncio.sleep(_EPSILON) + raise finally: task_done = True @@ -1145,6 +1148,34 @@ class BaseTaskTests: chained = cm.exception.__context__ self.assertEqual(type(chained), asyncio.CancelledError) + def test_wait_for_waits_for_task_cancellation_w_timeout_0(self): + loop = asyncio.new_event_loop() + self.addCleanup(loop.close) + + task_done = False + + async def foo(): + async def inner(): + nonlocal task_done + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + await asyncio.sleep(_EPSILON) + raise + finally: + task_done = True + + inner_task = self.new_task(loop, inner()) + await asyncio.sleep(_EPSILON) + await asyncio.wait_for(inner_task, timeout=0) + + with self.assertRaises(asyncio.TimeoutError) as cm: + loop.run_until_complete(foo()) + + self.assertTrue(task_done) + chained = cm.exception.__context__ + self.assertEqual(type(chained), asyncio.CancelledError) + def test_wait_for_reraises_exception_during_cancellation(self): loop = asyncio.new_event_loop() self.addCleanup(loop.close) -- cgit v1.2.1