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 at
   3  # https://github.com/aio-libs/async-timeout
   4  
   5  # It is vendored here to reduce chain-dependencies on this library, and
   6  # modified slightly to remove some features we don't use.
   7  
   8  
   9  import asyncio
  10  import warnings
  11  from types import TracebackType
  12  from typing import Any  # noqa
  13  from typing import Optional, Type
  14  
  15  
  16  class timeout:
  17      """timeout context manager.
  18  
  19      Useful in cases when you want to apply timeout logic around block
  20      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 logic
  28      loop - asyncio compatible event loop
  29      """
  30  
  31      def __init__(
  32          self,
  33          timeout: Optional[float],
  34          *,
  35          loop: Optional[asyncio.AbstractEventLoop] = None,
  36      ) -> None:
  37          self._timeout = timeout
  38          if loop is None:
  39              loop = asyncio.get_running_loop()
  40          else:
  41              warnings.warn(
  42                  """The loop argument to timeout() is deprecated.""", DeprecationWarning
  43              )
  44          self._loop = loop
  45          self._task = None  # type: Optional[asyncio.Task[Any]]
  46          self._cancelled = False
  47          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 None
  61  
  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      @property
  74      def expired(self) -> bool:
  75          return self._cancelled
  76  
  77      @property
  78      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 None
  83  
  84      def _do_enter(self) -> "timeout":
  85          # Support Tornado 5- without timeout
  86          # Details: https://github.com/python/asyncio/issues/392
  87          if self._timeout is None:
  88              return self
  89  
  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 self
  99  
 100          self._cancel_at = self._loop.time() + self._timeout
 101          self._cancel_handler = self._loop.call_at(self._cancel_at, self._cancel_task)
 102          return self
 103  
 104      def _do_exit(self, exc_type: Type[BaseException]) -> None:
 105          if exc_type is asyncio.CancelledError and self._cancelled:
 106              self._cancel_handler = None
 107              self._task = None
 108              raise asyncio.TimeoutError
 109          if self._timeout is not None and self._cancel_handler is not None:
 110              self._cancel_handler.cancel()
 111              self._cancel_handler = None
 112          self._task = None
 113          return None
 114  
 115      def _cancel_task(self) -> None:
 116          if self._task is not None:
 117              self._task.cancel()
 118              self._cancelled = True
 119