codekingpro/portable-devtools

1# This code is originally sourced from the aio-libs project "async_timeout",2# under the Apache 2.0 license. You may see the original project at3# https://github.com/aio-libs/async-timeout4 5# It is vendored here to reduce chain-dependencies on this library, and6# modified slightly to remove some features we don't use.7 8 9import asyncio10import warnings11from types import TracebackType12from typing import Any  # noqa13from typing import Optional, Type14 15 16class timeout:17    """timeout context manager.18 19    Useful in cases when you want to apply timeout logic around block20    of code or in cases when asyncio.wait_for is not suitable. For example:21 22    >>> with timeout(0.001):23    ...     async with aiohttp.get('https://github.com') as r:24    ...         await r.text()25 26 27    timeout - value in seconds or None to disable timeout logic28    loop - asyncio compatible event loop29    """30 31    def __init__(32        self,33        timeout: Optional[float],34        *,35        loop: Optional[asyncio.AbstractEventLoop] = None,36    ) -> None:37        self._timeout = timeout38        if loop is None:39            loop = asyncio.get_running_loop()40        else:41            warnings.warn(42                """The loop argument to timeout() is deprecated.""", DeprecationWarning43            )44        self._loop = loop45        self._task = None  # type: Optional[asyncio.Task[Any]]46        self._cancelled = False47        self._cancel_handler = None  # type: Optional[asyncio.Handle]48        self._cancel_at = None  # type: Optional[float]49 50    def __enter__(self) -> "timeout":51        return self._do_enter()52 53    def __exit__(54        self,55        exc_type: Type[BaseException],56        exc_val: BaseException,57        exc_tb: TracebackType,58    ) -> Optional[bool]:59        self._do_exit(exc_type)60        return None61 62    async def __aenter__(self) -> "timeout":63        return self._do_enter()64 65    async def __aexit__(66        self,67        exc_type: Type[BaseException],68        exc_val: BaseException,69        exc_tb: TracebackType,70    ) -> None:71        self._do_exit(exc_type)72 73    @property74    def expired(self) -> bool:75        return self._cancelled76 77    @property78    def remaining(self) -> Optional[float]:79        if self._cancel_at is not None:80            return max(self._cancel_at - self._loop.time(), 0.0)81        else:82            return None83 84    def _do_enter(self) -> "timeout":85        # Support Tornado 5- without timeout86        # Details: https://github.com/python/asyncio/issues/39287        if self._timeout is None:88            return self89 90        self._task = asyncio.current_task(self._loop)91        if self._task is None:92            raise RuntimeError(93                "Timeout context manager should be used " "inside a task"94            )95 96        if self._timeout <= 0:97            self._loop.call_soon(self._cancel_task)98            return self99 100        self._cancel_at = self._loop.time() + self._timeout101        self._cancel_handler = self._loop.call_at(self._cancel_at, self._cancel_task)102        return self103 104    def _do_exit(self, exc_type: Type[BaseException]) -> None:105        if exc_type is asyncio.CancelledError and self._cancelled:106            self._cancel_handler = None107            self._task = None108            raise asyncio.TimeoutError109        if self._timeout is not None and self._cancel_handler is not None:110            self._cancel_handler.cancel()111            self._cancel_handler = None112        self._task = None113        return None114 115    def _cancel_task(self) -> None:116        if self._task is not None:117            self._task.cancel()118            self._cancelled = True119