Sync with asyncio repo

1 files changed, 51 insertions, 0 deletions

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index a2ab8815b6..3c25e2d278 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -4,6 +4,7 @@ __all__ = ['Task',
            'FIRST_COMPLETED', 'FIRST_EXCEPTION', 'ALL_COMPLETED',
            'wait', 'wait_for', 'as_completed', 'sleep', 'async',
            'gather', 'shield', 'ensure_future', 'run_coroutine_threadsafe',
+           'timeout',
            ]
 
 import concurrent.futures
@@ -732,3 +733,53 @@ def run_coroutine_threadsafe(coro, loop):
 
     loop.call_soon_threadsafe(callback)
     return future
+
+
+def timeout(timeout, *, loop=None):
+    """A factory which produce a context manager with timeout.
+
+    Useful in cases when you want to apply timeout logic around block
+    of code or in cases when asyncio.wait_for is not suitable.
+
+    For example:
+
+    >>> with asyncio.timeout(0.001):
+    >>>     yield from coro()
+
+
+    timeout: timeout value in seconds
+    loop: asyncio compatible event loop
+    """
+    if loop is None:
+        loop = events.get_event_loop()
+    return _Timeout(timeout, loop=loop)
+
+
+class _Timeout:
+    def __init__(self, timeout, *, loop):
+        self._timeout = timeout
+        self._loop = loop
+        self._task = None
+        self._cancelled = False
+        self._cancel_handler = None
+
+    def __enter__(self):
+        self._task = Task.current_task(loop=self._loop)
+        if self._task is None:
+            raise RuntimeError('Timeout context manager should be used '
+                               'inside a task')
+        self._cancel_handler = self._loop.call_later(
+            self._timeout, self._cancel_task)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type is futures.CancelledError and self._cancelled:
+            self._cancel_handler = None
+            self._task = None
+            raise futures.TimeoutError
+        self._cancel_handler.cancel()
+        self._cancel_handler = None
+        self._task = None
+
+    def _cancel_task(self):
+        self._cancelled = self._task.cancel()