rewrite wait_for to use timeouts.timeout

Author: kumaraditya303

Lib/asyncio/tasks.py

@@ -24,6 +24,7 @@
 from . import events
 from . import exceptions
 from . import futures
+from . import timeouts
 from .coroutines import _is_coroutine
 
 # Helper to generate new task names
@@ -439,63 +440,23 @@ async def wait_for(fut, timeout):
 
     This function is a coroutine.
     """
-    loop = events.get_running_loop()
-
     if timeout is None:
         return await fut
 
     if timeout <= 0:
-        fut = ensure_future(fut, loop=loop)
+        fut = ensure_future(fut)
 
         if fut.done():
             return fut.result()
 
-        await _cancel_and_wait(fut, loop=loop)
+        await _cancel_and_wait(fut)
         try:
             return fut.result()
         except exceptions.CancelledError as exc:
-            raise exceptions.TimeoutError() from exc
-
-    waiter = loop.create_future()
-    timeout_handle = loop.call_later(timeout, _release_waiter, waiter)
-    cb = functools.partial(_release_waiter, waiter)
-
-    fut = ensure_future(fut, loop=loop)
-    fut.add_done_callback(cb)
-
-    try:
-        # wait until the future completes or the timeout
-        try:
-            await waiter
-        except exceptions.CancelledError:
-            if fut.done():
-                return fut.result()
-            else:
-                fut.remove_done_callback(cb)
-                # We must ensure that the task is not running
-                # after wait_for() returns.
-                # See https://bugs.python.org/issue32751
-                await _cancel_and_wait(fut, loop=loop)
-                raise
-
-        if fut.done():
-            return fut.result()
-        else:
-            fut.remove_done_callback(cb)
-            # We must ensure that the task is not running
-            # after wait_for() returns.
-            # See https://bugs.python.org/issue32751
-            await _cancel_and_wait(fut, loop=loop)
-            # In case task cancellation failed with some
-            # exception, we should re-raise it
-            # See https://bugs.python.org/issue40607
-            try:
-                return fut.result()
-            except exceptions.CancelledError as exc:
-                raise exceptions.TimeoutError() from exc
-    finally:
-        timeout_handle.cancel()
+            raise TimeoutError from exc
 
+    async with timeouts.timeout(timeout):
+        return await fut
 
 async def _wait(fs, timeout, return_when, loop):
     """Internal helper for wait().
@@ -541,9 +502,10 @@ def _on_completion(f):
     return done, pending
 
 
-async def _cancel_and_wait(fut, loop):
+async def _cancel_and_wait(fut):
     """Cancel the *fut* future or task and wait until it completes."""
 
+    loop = events.get_running_loop()
     waiter = loop.create_future()
     cb = functools.partial(_release_waiter, waiter)
     fut.add_done_callback(cb)

Lib/test/test_asyncio/test_waitfor.py

@@ -237,33 +237,6 @@ async def inner():
         with self.assertRaises(FooException):
             await foo()
 
-    async def test_wait_for_self_cancellation(self):
-        async def inner():
-            try:
-                await asyncio.sleep(0.3)
-            except asyncio.CancelledError:
-                try:
-                    await asyncio.sleep(0.3)
-                except asyncio.CancelledError:
-                    await asyncio.sleep(0.3)
-
-            return 42
-
-        inner_task = asyncio.create_task(inner())
-
-        wait = asyncio.wait_for(inner_task, timeout=0.1)
-
-        # Test that wait_for itself is properly cancellable
-        # even when the initial task holds up the initial cancellation.
-        task = asyncio.create_task(wait)
-        await asyncio.sleep(0.2)
-        task.cancel()
-
-        with self.assertRaises(asyncio.CancelledError):
-            await task
-
-        self.assertEqual(await inner_task, 42)
-
     async def _test_cancel_wait_for(self, timeout):
         loop = asyncio.get_running_loop()