async_timeout.py 8.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282
  1. # From https://github.com/aio-libs/async-timeout/blob/master/async_timeout/__init__.py
  2. # Licensed under the Apache License (Apache-2.0)
  3. import asyncio
  4. import enum
  5. import sys
  6. import warnings
  7. from types import TracebackType
  8. from typing import Optional, Type
  9. if sys.version_info >= (3, 11):
  10. from typing import final
  11. else:
  12. # From https://github.com/python/typing_extensions/blob/main/src/typing_extensions.py
  13. # Licensed under the Python Software Foundation License (PSF-2.0)
  14. # @final exists in 3.8+, but we backport it for all versions
  15. # before 3.11 to keep support for the __final__ attribute.
  16. # See https://bugs.python.org/issue46342
  17. def final(f):
  18. """This decorator can be used to indicate to type checkers that
  19. the decorated method cannot be overridden, and decorated class
  20. cannot be subclassed. For example:
  21. class Base:
  22. @final
  23. def done(self) -> None:
  24. ...
  25. class Sub(Base):
  26. def done(self) -> None: # Error reported by type checker
  27. ...
  28. @final
  29. class Leaf:
  30. ...
  31. class Other(Leaf): # Error reported by type checker
  32. ...
  33. There is no runtime checking of these properties. The decorator
  34. sets the ``__final__`` attribute to ``True`` on the decorated object
  35. to allow runtime introspection.
  36. """
  37. try:
  38. f.__final__ = True
  39. except (AttributeError, TypeError):
  40. # Skip the attribute silently if it is not writable.
  41. # AttributeError happens if the object has __slots__ or a
  42. # read-only property, TypeError if it's a builtin class.
  43. pass
  44. return f
  45. # End https://github.com/python/typing_extensions/blob/main/src/typing_extensions.py
  46. if sys.version_info >= (3, 11):
  47. def _uncancel_task(task: "asyncio.Task[object]") -> None:
  48. task.uncancel()
  49. else:
  50. def _uncancel_task(task: "asyncio.Task[object]") -> None:
  51. pass
  52. __version__ = "4.0.3"
  53. __all__ = ("timeout", "timeout_at", "Timeout")
  54. def timeout(delay: Optional[float]) -> "Timeout":
  55. """timeout context manager.
  56. Useful in cases when you want to apply timeout logic around block
  57. of code or in cases when asyncio.wait_for is not suitable. For example:
  58. >>> async with timeout(0.001):
  59. ... async with aiohttp.get('https://github.com') as r:
  60. ... await r.text()
  61. delay - value in seconds or None to disable timeout logic
  62. """
  63. loop = asyncio.get_running_loop()
  64. if delay is not None:
  65. deadline = loop.time() + delay # type: Optional[float]
  66. else:
  67. deadline = None
  68. return Timeout(deadline, loop)
  69. def timeout_at(deadline: Optional[float]) -> "Timeout":
  70. """Schedule the timeout at absolute time.
  71. deadline argument points on the time in the same clock system
  72. as loop.time().
  73. Please note: it is not POSIX time but a time with
  74. undefined starting base, e.g. the time of the system power on.
  75. >>> async with timeout_at(loop.time() + 10):
  76. ... async with aiohttp.get('https://github.com') as r:
  77. ... await r.text()
  78. """
  79. loop = asyncio.get_running_loop()
  80. return Timeout(deadline, loop)
  81. class _State(enum.Enum):
  82. INIT = "INIT"
  83. ENTER = "ENTER"
  84. TIMEOUT = "TIMEOUT"
  85. EXIT = "EXIT"
  86. @final
  87. class Timeout:
  88. # Internal class, please don't instantiate it directly
  89. # Use timeout() and timeout_at() public factories instead.
  90. #
  91. # Implementation note: `async with timeout()` is preferred
  92. # over `with timeout()`.
  93. # While technically the Timeout class implementation
  94. # doesn't need to be async at all,
  95. # the `async with` statement explicitly points that
  96. # the context manager should be used from async function context.
  97. #
  98. # This design allows to avoid many silly misusages.
  99. #
  100. # TimeoutError is raised immediately when scheduled
  101. # if the deadline is passed.
  102. # The purpose is to time out as soon as possible
  103. # without waiting for the next await expression.
  104. __slots__ = ("_deadline", "_loop", "_state", "_timeout_handler", "_task")
  105. def __init__(
  106. self, deadline: Optional[float], loop: asyncio.AbstractEventLoop
  107. ) -> None:
  108. self._loop = loop
  109. self._state = _State.INIT
  110. self._task: Optional["asyncio.Task[object]"] = None
  111. self._timeout_handler = None # type: Optional[asyncio.Handle]
  112. if deadline is None:
  113. self._deadline = None # type: Optional[float]
  114. else:
  115. self.update(deadline)
  116. def __enter__(self) -> "Timeout":
  117. warnings.warn(
  118. "with timeout() is deprecated, use async with timeout() instead",
  119. DeprecationWarning,
  120. stacklevel=2,
  121. )
  122. self._do_enter()
  123. return self
  124. def __exit__(
  125. self,
  126. exc_type: Optional[Type[BaseException]],
  127. exc_val: Optional[BaseException],
  128. exc_tb: Optional[TracebackType],
  129. ) -> Optional[bool]:
  130. self._do_exit(exc_type)
  131. return None
  132. async def __aenter__(self) -> "Timeout":
  133. self._do_enter()
  134. return self
  135. async def __aexit__(
  136. self,
  137. exc_type: Optional[Type[BaseException]],
  138. exc_val: Optional[BaseException],
  139. exc_tb: Optional[TracebackType],
  140. ) -> Optional[bool]:
  141. self._do_exit(exc_type)
  142. return None
  143. @property
  144. def expired(self) -> bool:
  145. """Is timeout expired during execution?"""
  146. return self._state == _State.TIMEOUT
  147. @property
  148. def deadline(self) -> Optional[float]:
  149. return self._deadline
  150. def reject(self) -> None:
  151. """Reject scheduled timeout if any."""
  152. # cancel is maybe better name but
  153. # task.cancel() raises CancelledError in asyncio world.
  154. if self._state not in (_State.INIT, _State.ENTER):
  155. raise RuntimeError(f"invalid state {self._state.value}")
  156. self._reject()
  157. def _reject(self) -> None:
  158. self._task = None
  159. if self._timeout_handler is not None:
  160. self._timeout_handler.cancel()
  161. self._timeout_handler = None
  162. def shift(self, delay: float) -> None:
  163. """Advance timeout on delay seconds.
  164. The delay can be negative.
  165. Raise RuntimeError if shift is called when deadline is not scheduled
  166. """
  167. deadline = self._deadline
  168. if deadline is None:
  169. raise RuntimeError("cannot shift timeout if deadline is not scheduled")
  170. self.update(deadline + delay)
  171. def update(self, deadline: float) -> None:
  172. """Set deadline to absolute value.
  173. deadline argument points on the time in the same clock system
  174. as loop.time().
  175. If new deadline is in the past the timeout is raised immediately.
  176. Please note: it is not POSIX time but a time with
  177. undefined starting base, e.g. the time of the system power on.
  178. """
  179. if self._state == _State.EXIT:
  180. raise RuntimeError("cannot reschedule after exit from context manager")
  181. if self._state == _State.TIMEOUT:
  182. raise RuntimeError("cannot reschedule expired timeout")
  183. if self._timeout_handler is not None:
  184. self._timeout_handler.cancel()
  185. self._deadline = deadline
  186. if self._state != _State.INIT:
  187. self._reschedule()
  188. def _reschedule(self) -> None:
  189. assert self._state == _State.ENTER
  190. deadline = self._deadline
  191. if deadline is None:
  192. return
  193. now = self._loop.time()
  194. if self._timeout_handler is not None:
  195. self._timeout_handler.cancel()
  196. self._task = asyncio.current_task()
  197. if deadline <= now:
  198. self._timeout_handler = self._loop.call_soon(self._on_timeout)
  199. else:
  200. self._timeout_handler = self._loop.call_at(deadline, self._on_timeout)
  201. def _do_enter(self) -> None:
  202. if self._state != _State.INIT:
  203. raise RuntimeError(f"invalid state {self._state.value}")
  204. self._state = _State.ENTER
  205. self._reschedule()
  206. def _do_exit(self, exc_type: Optional[Type[BaseException]]) -> None:
  207. if exc_type is asyncio.CancelledError and self._state == _State.TIMEOUT:
  208. assert self._task is not None
  209. _uncancel_task(self._task)
  210. self._timeout_handler = None
  211. self._task = None
  212. raise asyncio.TimeoutError
  213. # timeout has not expired
  214. self._state = _State.EXIT
  215. self._reject()
  216. return None
  217. def _on_timeout(self) -> None:
  218. assert self._task is not None
  219. self._task.cancel()
  220. self._state = _State.TIMEOUT
  221. # drop the reference early
  222. self._timeout_handler = None
  223. # End https://github.com/aio-libs/async-timeout/blob/master/async_timeout/__init__.py