matteo-the-prestige/matteo_env/Lib/site-packages/async_timeout/__init__.py
2020-12-19 19:08:09 -05:00

116 lines
3.6 KiB
Python

import asyncio
import sys
from types import TracebackType
from typing import Optional, Type, Any # noqa
__version__ = '3.0.1'
PY_37 = sys.version_info >= (3, 7)
class timeout:
"""timeout context manager.
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 timeout(0.001):
... async with aiohttp.get('https://github.com') as r:
... await r.text()
timeout - value in seconds or None to disable timeout logic
loop - asyncio compatible event loop
"""
def __init__(self, timeout: Optional[float],
*, loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
self._timeout = timeout
if loop is None:
loop = asyncio.get_event_loop()
self._loop = loop
self._task = None # type: Optional[asyncio.Task[Any]]
self._cancelled = False
self._cancel_handler = None # type: Optional[asyncio.Handle]
self._cancel_at = None # type: Optional[float]
def __enter__(self) -> 'timeout':
return self._do_enter()
def __exit__(self,
exc_type: Type[BaseException],
exc_val: BaseException,
exc_tb: TracebackType) -> Optional[bool]:
self._do_exit(exc_type)
return None
async def __aenter__(self) -> 'timeout':
return self._do_enter()
async def __aexit__(self,
exc_type: Type[BaseException],
exc_val: BaseException,
exc_tb: TracebackType) -> None:
self._do_exit(exc_type)
@property
def expired(self) -> bool:
return self._cancelled
@property
def remaining(self) -> Optional[float]:
if self._cancel_at is not None:
return max(self._cancel_at - self._loop.time(), 0.0)
else:
return None
def _do_enter(self) -> 'timeout':
# Support Tornado 5- without timeout
# Details: https://github.com/python/asyncio/issues/392
if self._timeout is None:
return self
self._task = current_task(self._loop)
if self._task is None:
raise RuntimeError('Timeout context manager should be used '
'inside a task')
if self._timeout <= 0:
self._loop.call_soon(self._cancel_task)
return self
self._cancel_at = self._loop.time() + self._timeout
self._cancel_handler = self._loop.call_at(
self._cancel_at, self._cancel_task)
return self
def _do_exit(self, exc_type: Type[BaseException]) -> None:
if exc_type is asyncio.CancelledError and self._cancelled:
self._cancel_handler = None
self._task = None
raise asyncio.TimeoutError
if self._timeout is not None and self._cancel_handler is not None:
self._cancel_handler.cancel()
self._cancel_handler = None
self._task = None
return None
def _cancel_task(self) -> None:
if self._task is not None:
self._task.cancel()
self._cancelled = True
def current_task(loop: asyncio.AbstractEventLoop) -> 'asyncio.Task[Any]':
if PY_37:
task = asyncio.current_task(loop=loop) # type: ignore
else:
task = asyncio.Task.current_task(loop=loop)
if task is None:
# this should be removed, tokio must use register_task and family API
if hasattr(loop, 'current_task'):
task = loop.current_task() # type: ignore
return task