Sync with asyncio repo
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index a2ab881..3c25e2d 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -4,6 +4,7 @@
'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 @@
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()