bpo-32751: Wait for task cancel in asyncio.wait_for() when timeout <= 0 (#21895)

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.

1 files changed, 7 insertions, 2 deletions

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index c37f0e1..7ecec96 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)