__init__.py 9.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276
  1. import asyncio
  2. import enum
  3. import sys
  4. from types import TracebackType
  5. from typing import Optional, Type, final
  6. __version__ = "5.0.1"
  7. __all__ = ("timeout", "timeout_at", "Timeout")
  8. def timeout(delay: Optional[float]) -> "Timeout":
  9. """timeout context manager.
  10. Useful in cases when you want to apply timeout logic around block
  11. of code or in cases when asyncio.wait_for is not suitable. For example:
  12. >>> async with timeout(0.001):
  13. ... async with aiohttp.get('https://github.com') as r:
  14. ... await r.text()
  15. delay - value in seconds or None to disable timeout logic
  16. """
  17. loop = asyncio.get_running_loop()
  18. if delay is not None:
  19. deadline = loop.time() + delay # type: Optional[float]
  20. else:
  21. deadline = None
  22. return Timeout(deadline, loop)
  23. def timeout_at(deadline: Optional[float]) -> "Timeout":
  24. """Schedule the timeout at absolute time.
  25. deadline argument points on the time in the same clock system
  26. as loop.time().
  27. Please note: it is not POSIX time but a time with
  28. undefined starting base, e.g. the time of the system power on.
  29. >>> async with timeout_at(loop.time() + 10):
  30. ... async with aiohttp.get('https://github.com') as r:
  31. ... await r.text()
  32. """
  33. loop = asyncio.get_running_loop()
  34. return Timeout(deadline, loop)
  35. class _State(enum.Enum):
  36. INIT = "INIT"
  37. ENTER = "ENTER"
  38. TIMEOUT = "TIMEOUT"
  39. EXIT = "EXIT"
  40. if sys.version_info >= (3, 11):
  41. class _Expired:
  42. __slots__ = ("_val",)
  43. def __init__(self, val: bool) -> None:
  44. self._val = val
  45. def __call__(self) -> bool:
  46. return self._val
  47. def __bool__(self) -> bool:
  48. return self._val
  49. def __repr__(self) -> str:
  50. return repr(self._val)
  51. def __str__(self) -> str:
  52. return str(self._val)
  53. @final
  54. class Timeout(asyncio.Timeout): # type: ignore[misc]
  55. # Supports full asyncio.Timeout API.
  56. # Also provides several asyncio_timeout specific methods
  57. # for backward compatibility.
  58. def __init__(
  59. self, deadline: Optional[float], loop: asyncio.AbstractEventLoop
  60. ) -> None:
  61. super().__init__(deadline)
  62. @property
  63. def expired(self) -> _Expired:
  64. # a hacky property hat can provide both roles:
  65. # timeout.expired() from asyncio
  66. # timeout.expired from asyncio_timeout
  67. return _Expired(super().expired())
  68. @property
  69. def deadline(self) -> Optional[float]:
  70. return self.when()
  71. def reject(self) -> None:
  72. """Reject scheduled timeout if any."""
  73. # cancel is maybe better name but
  74. # task.cancel() raises CancelledError in asyncio world.
  75. self.reschedule(None)
  76. def shift(self, delay: float) -> None:
  77. """Advance timeout on delay seconds.
  78. The delay can be negative.
  79. Raise RuntimeError if shift is called when deadline is not scheduled
  80. """
  81. deadline = self.when()
  82. if deadline is None:
  83. raise RuntimeError("cannot shift timeout if deadline is not scheduled")
  84. self.reschedule(deadline + delay)
  85. def update(self, deadline: float) -> None:
  86. """Set deadline to absolute value.
  87. deadline argument points on the time in the same clock system
  88. as loop.time().
  89. If new deadline is in the past the timeout is raised immediately.
  90. Please note: it is not POSIX time but a time with
  91. undefined starting base, e.g. the time of the system power on.
  92. """
  93. self.reschedule(deadline)
  94. else:
  95. @final
  96. class Timeout:
  97. # Internal class, please don't instantiate it directly
  98. # Use timeout() and timeout_at() public factories instead.
  99. #
  100. # Implementation note: `async with timeout()` is preferred
  101. # over `with timeout()`.
  102. # While technically the Timeout class implementation
  103. # doesn't need to be async at all,
  104. # the `async with` statement explicitly points that
  105. # the context manager should be used from async function context.
  106. #
  107. # This design allows to avoid many silly misusages.
  108. #
  109. # TimeoutError is raised immediately when scheduled
  110. # if the deadline is passed.
  111. # The purpose is to time out as soon as possible
  112. # without waiting for the next await expression.
  113. __slots__ = ("_deadline", "_loop", "_state", "_timeout_handler", "_task")
  114. def __init__(
  115. self, deadline: Optional[float], loop: asyncio.AbstractEventLoop
  116. ) -> None:
  117. self._loop = loop
  118. self._state = _State.INIT
  119. self._task: Optional["asyncio.Task[object]"] = None
  120. self._timeout_handler = None # type: Optional[asyncio.Handle]
  121. if deadline is None:
  122. self._deadline = None # type: Optional[float]
  123. else:
  124. self.update(deadline)
  125. async def __aenter__(self) -> "Timeout":
  126. self._do_enter()
  127. return self
  128. async def __aexit__(
  129. self,
  130. exc_type: Optional[Type[BaseException]],
  131. exc_val: Optional[BaseException],
  132. exc_tb: Optional[TracebackType],
  133. ) -> Optional[bool]:
  134. self._do_exit(exc_type)
  135. return None
  136. @property
  137. def expired(self) -> bool:
  138. """Is timeout expired during execution?"""
  139. return self._state == _State.TIMEOUT
  140. @property
  141. def deadline(self) -> Optional[float]:
  142. return self._deadline
  143. def reject(self) -> None:
  144. """Reject scheduled timeout if any."""
  145. # cancel is maybe better name but
  146. # task.cancel() raises CancelledError in asyncio world.
  147. if self._state not in (_State.INIT, _State.ENTER):
  148. raise RuntimeError(f"invalid state {self._state.value}")
  149. self._reject()
  150. def _reject(self) -> None:
  151. self._task = None
  152. if self._timeout_handler is not None:
  153. self._timeout_handler.cancel()
  154. self._timeout_handler = None
  155. def shift(self, delay: float) -> None:
  156. """Advance timeout on delay seconds.
  157. The delay can be negative.
  158. Raise RuntimeError if shift is called when deadline is not scheduled
  159. """
  160. deadline = self._deadline
  161. if deadline is None:
  162. raise RuntimeError("cannot shift timeout if deadline is not scheduled")
  163. self.update(deadline + delay)
  164. def update(self, deadline: float) -> None:
  165. """Set deadline to absolute value.
  166. deadline argument points on the time in the same clock system
  167. as loop.time().
  168. If new deadline is in the past the timeout is raised immediately.
  169. Please note: it is not POSIX time but a time with
  170. undefined starting base, e.g. the time of the system power on.
  171. """
  172. if self._state == _State.EXIT:
  173. raise RuntimeError("cannot reschedule after exit from context manager")
  174. if self._state == _State.TIMEOUT:
  175. raise RuntimeError("cannot reschedule expired timeout")
  176. if self._timeout_handler is not None:
  177. self._timeout_handler.cancel()
  178. self._deadline = deadline
  179. if self._state != _State.INIT:
  180. self._reschedule()
  181. def _reschedule(self) -> None:
  182. assert self._state == _State.ENTER
  183. deadline = self._deadline
  184. if deadline is None:
  185. return
  186. now = self._loop.time()
  187. if self._timeout_handler is not None:
  188. self._timeout_handler.cancel()
  189. self._task = asyncio.current_task()
  190. if deadline <= now:
  191. self._timeout_handler = self._loop.call_soon(self._on_timeout)
  192. else:
  193. self._timeout_handler = self._loop.call_at(deadline, self._on_timeout)
  194. def _do_enter(self) -> None:
  195. if self._state != _State.INIT:
  196. raise RuntimeError(f"invalid state {self._state.value}")
  197. self._state = _State.ENTER
  198. self._reschedule()
  199. def _do_exit(self, exc_type: Optional[Type[BaseException]]) -> None:
  200. if exc_type is asyncio.CancelledError and self._state == _State.TIMEOUT:
  201. assert self._task is not None
  202. self._timeout_handler = None
  203. self._task = None
  204. raise asyncio.TimeoutError
  205. # timeout has not expired
  206. self._state = _State.EXIT
  207. self._reject()
  208. return None
  209. def _on_timeout(self) -> None:
  210. assert self._task is not None
  211. self._task.cancel()
  212. self._state = _State.TIMEOUT
  213. # drop the reference early
  214. self._timeout_handler = None