|
|
- 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
|