_tasks.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412
  1. from __future__ import annotations
  2. import math
  3. import sys
  4. from collections.abc import (
  5. Coroutine,
  6. Generator,
  7. )
  8. from contextlib import (
  9. contextmanager,
  10. )
  11. from enum import Enum, auto
  12. from inspect import iscoroutine
  13. from types import TracebackType
  14. from typing import Any, Generic, final
  15. from ..abc import TaskGroup, TaskStatus
  16. from ._eventloop import get_async_backend, get_cancelled_exc_class
  17. from ._exceptions import TaskCancelled, TaskFailed, TaskNotFinished
  18. if sys.version_info >= (3, 13):
  19. from typing import TypeVar
  20. else:
  21. from typing_extensions import TypeVar
  22. if sys.version_info >= (3, 11):
  23. from typing import Never, TypeVarTuple
  24. else:
  25. from typing_extensions import Never, TypeVarTuple
  26. T = TypeVar("T")
  27. T_co = TypeVar("T_co", covariant=True)
  28. T_startval = TypeVar("T_startval", covariant=True, default=Never)
  29. PosArgsT = TypeVarTuple("PosArgsT")
  30. class _IgnoredTaskStatus(TaskStatus[object]):
  31. def started(self, value: object = None) -> None:
  32. pass
  33. TASK_STATUS_IGNORED = _IgnoredTaskStatus()
  34. class CancelScope:
  35. """
  36. Wraps a unit of work that can be made separately cancellable.
  37. :param deadline: The time (clock value) when this scope is cancelled automatically
  38. :param shield: ``True`` to shield the cancel scope from external cancellation
  39. :raises NoEventLoopError: if no supported asynchronous event loop is running in the
  40. current thread
  41. """
  42. __slots__ = ("__weakref__",)
  43. def __new__(
  44. cls, *, deadline: float = math.inf, shield: bool = False
  45. ) -> CancelScope:
  46. return get_async_backend().create_cancel_scope(shield=shield, deadline=deadline)
  47. def cancel(self, reason: str | None = None) -> None:
  48. """
  49. Cancel this scope immediately.
  50. :param reason: a message describing the reason for the cancellation
  51. """
  52. raise NotImplementedError
  53. @property
  54. def deadline(self) -> float:
  55. """
  56. The time (clock value) when this scope is cancelled automatically.
  57. Will be ``float('inf')`` if no timeout has been set.
  58. """
  59. raise NotImplementedError
  60. @deadline.setter
  61. def deadline(self, value: float) -> None:
  62. raise NotImplementedError
  63. @property
  64. def cancel_called(self) -> bool:
  65. """``True`` if :meth:`cancel` has been called."""
  66. raise NotImplementedError
  67. @property
  68. def cancelled_caught(self) -> bool:
  69. """
  70. ``True`` if this scope suppressed a cancellation exception it itself raised.
  71. This is typically used to check if any work was interrupted, or to see if the
  72. scope was cancelled due to its deadline being reached. The value will, however,
  73. only be ``True`` if the cancellation was triggered by the scope itself (and not
  74. an outer scope).
  75. """
  76. raise NotImplementedError
  77. @property
  78. def shield(self) -> bool:
  79. """
  80. ``True`` if this scope is shielded from external cancellation.
  81. While a scope is shielded, it will not receive cancellations from outside.
  82. """
  83. raise NotImplementedError
  84. @shield.setter
  85. def shield(self, value: bool) -> None:
  86. raise NotImplementedError
  87. def __enter__(self) -> CancelScope:
  88. raise NotImplementedError
  89. def __exit__(
  90. self,
  91. exc_type: type[BaseException] | None,
  92. exc_val: BaseException | None,
  93. exc_tb: TracebackType | None,
  94. ) -> bool:
  95. raise NotImplementedError
  96. @contextmanager
  97. def fail_after(
  98. delay: float | None, shield: bool = False
  99. ) -> Generator[CancelScope, None, None]:
  100. """
  101. Create a context manager which raises a :class:`TimeoutError` if does not finish in
  102. time.
  103. :param delay: maximum allowed time (in seconds) before raising the exception, or
  104. ``None`` to disable the timeout
  105. :param shield: ``True`` to shield the cancel scope from external cancellation
  106. :return: a context manager that yields a cancel scope
  107. :rtype: :class:`~typing.ContextManager`\\[:class:`~anyio.CancelScope`\\]
  108. :raises NoEventLoopError: if no supported asynchronous event loop is running in the
  109. current thread
  110. """
  111. current_time = get_async_backend().current_time
  112. deadline = (current_time() + delay) if delay is not None else math.inf
  113. with get_async_backend().create_cancel_scope(
  114. deadline=deadline, shield=shield
  115. ) as cancel_scope:
  116. yield cancel_scope
  117. if cancel_scope.cancelled_caught and current_time() >= cancel_scope.deadline:
  118. raise TimeoutError
  119. def move_on_after(delay: float | None, shield: bool = False) -> CancelScope:
  120. """
  121. Create a cancel scope with a deadline that expires after the given delay.
  122. :param delay: maximum allowed time (in seconds) before exiting the context block, or
  123. ``None`` to disable the timeout
  124. :param shield: ``True`` to shield the cancel scope from external cancellation
  125. :return: a cancel scope
  126. :raises NoEventLoopError: if no supported asynchronous event loop is running in the
  127. current thread
  128. """
  129. deadline = (
  130. (get_async_backend().current_time() + delay) if delay is not None else math.inf
  131. )
  132. return get_async_backend().create_cancel_scope(deadline=deadline, shield=shield)
  133. def current_effective_deadline() -> float:
  134. """
  135. Return the nearest deadline among all the cancel scopes effective for the current
  136. task.
  137. :return: a clock value from the event loop's internal clock (or ``float('inf')`` if
  138. there is no deadline in effect, or ``float('-inf')`` if the current scope has
  139. been cancelled)
  140. :rtype: float
  141. :raises NoEventLoopError: if no supported asynchronous event loop is running in the
  142. current thread
  143. """
  144. return get_async_backend().current_effective_deadline()
  145. def create_task_group() -> TaskGroup:
  146. """
  147. Create a task group.
  148. :return: a task group
  149. :raises NoEventLoopError: if no supported asynchronous event loop is running in the
  150. current thread
  151. """
  152. return get_async_backend().create_task_group()
  153. @final
  154. class TaskHandle(Generic[T_co, T_startval]):
  155. """
  156. Returned from the task-spawning methods of :class:`TaskGroup`. Can be awaited on to
  157. get the return value of the task (or the raised exception). If the task was
  158. terminated by a :exc:`BaseException`, :exc:`TaskFailed` will be raised (or its
  159. subclass :exc:`TaskCancelled` if the task was cancelled).
  160. .. versionadded:: 4.14.0
  161. """
  162. class Status(Enum):
  163. """
  164. The status of a task handle.
  165. .. attribute:: PENDING
  166. The task has not finished yet.
  167. .. attribute:: FINISHED
  168. The task has finished with a return value.
  169. .. attribute:: CANCELLING
  170. The task has been cancelled but has not finished yet.
  171. .. attribute:: CANCELLED
  172. The task was cancelled and has finished since.
  173. .. attribute:: FAILED
  174. The task raised an exception.
  175. """
  176. PENDING = auto()
  177. FINISHED = auto()
  178. CANCELLING = auto()
  179. CANCELLED = auto()
  180. FAILED = auto()
  181. __slots__ = (
  182. "__weakref__",
  183. "_coro",
  184. "_name",
  185. "_cancel_scope",
  186. "_finished_event",
  187. "_return_value",
  188. "_start_value",
  189. "_exception",
  190. )
  191. _return_value: T_co
  192. _start_value: T_startval
  193. def __init__(self, coro: Coroutine[Any, Any, T_co], name: object) -> None:
  194. from ._synchronization import Event
  195. self._coro = coro
  196. self._cancel_scope = CancelScope()
  197. self._finished_event = Event()
  198. self._exception: BaseException | None = None
  199. if name is not None:
  200. self._name = str(name)
  201. elif iscoroutine(coro):
  202. self._name = coro.__qualname__
  203. else:
  204. self._name = str(coro) # coroutine-like object (e.g. asend() objects)
  205. async def _run_coro(self) -> None:
  206. __tracebackhide__ = True
  207. with self._cancel_scope:
  208. try:
  209. retval = await self._coro
  210. except BaseException as exc:
  211. self._exception = exc
  212. raise
  213. else:
  214. self._return_value = retval
  215. finally:
  216. self._finished_event.set()
  217. del self # Break the reference cycle
  218. def cancel(self) -> None:
  219. """
  220. Set the task to a cancelled state.
  221. This will interrupt any interruptible asynchronous operation, and will cause
  222. any further awaits on this task to get immediately cancelled, unless done in
  223. a shielded cancel scope.
  224. If the task has already finished, this method has no effect.
  225. """
  226. if not self._finished_event.is_set():
  227. self._cancel_scope.cancel()
  228. @property
  229. def coro(self) -> Coroutine[Any, Any, T_co]:
  230. """
  231. The coroutine object that was passed to one of the task-spawning methods in
  232. :class:`TaskGroup`.
  233. """
  234. return self._coro
  235. @property
  236. def status(self) -> TaskHandle.Status:
  237. """
  238. The current status of the task.
  239. Every task starts in the :attr:`~TaskHandle.Status.PENDING` state.
  240. If a task is cancelled while in this state, it will transition to the
  241. :attr:`~TaskHandle.Status.CANCELLING` state. When the task finishes, it will
  242. transition to one of the three final states (
  243. :attr:`~TaskHandle.Status.FINISHED`, :attr:`~TaskHandle.Status.FAILED`, or
  244. :attr:`~TaskHandle.Status.CANCELLING`) depending on the exception the task
  245. raised, if any. No other status transitions will happen.
  246. """
  247. if not self._finished_event.is_set():
  248. if self._cancel_scope.cancel_called:
  249. return TaskHandle.Status.CANCELLING
  250. else:
  251. return TaskHandle.Status.PENDING
  252. elif self._exception is not None:
  253. if isinstance(self._exception, get_cancelled_exc_class()):
  254. return TaskHandle.Status.CANCELLED
  255. else:
  256. return TaskHandle.Status.FAILED
  257. else:
  258. return TaskHandle.Status.FINISHED
  259. @property
  260. def name(self) -> str:
  261. """The name of the task."""
  262. return self._name
  263. @property
  264. def exception(self) -> BaseException | None:
  265. """
  266. The exception raised by the task, or ``None`` if it finished without raising.
  267. :raises TaskNotFinished: if the task has not finished yet
  268. :raises TaskCancelled: if the task was cancelled
  269. """
  270. match self.status:
  271. case TaskHandle.Status.PENDING:
  272. raise TaskNotFinished("the task has not finished yet")
  273. case TaskHandle.Status.FINISHED:
  274. return None
  275. case TaskHandle.Status.CANCELLING:
  276. raise TaskCancelled("the task was cancelled")
  277. case TaskHandle.Status.CANCELLED:
  278. raise TaskCancelled("the task was cancelled") from self._exception
  279. case TaskHandle.Status.FAILED:
  280. return self._exception
  281. @property
  282. def return_value(self) -> T_co:
  283. """
  284. The return value of the task.
  285. :raises TaskNotFinished: if the task has not finished yet
  286. :raises TaskCancelled: if the task was cancelled
  287. :raises TaskFailed: if the task raised an exception
  288. """
  289. match self.status:
  290. case TaskHandle.Status.PENDING:
  291. raise TaskNotFinished("the task has not finished yet")
  292. case TaskHandle.Status.FINISHED:
  293. return self._return_value
  294. case TaskHandle.Status.CANCELLING:
  295. raise TaskCancelled("the task was cancelled")
  296. case TaskHandle.Status.CANCELLED:
  297. raise TaskCancelled("the task was cancelled") from self._exception
  298. case TaskHandle.Status.FAILED:
  299. raise TaskFailed("the task raised an exception") from self._exception
  300. @property
  301. def start_value(self) -> T_startval:
  302. """
  303. The value passed to :meth:`task_status.started() <.abc.TaskStatus.started>`,
  304. :raises RuntimeError: if the task was not started with :meth:`TaskGroup.start()
  305. <.abc.TaskGroup.start>`
  306. """
  307. try:
  308. return self._start_value
  309. except AttributeError:
  310. raise RuntimeError(
  311. "the task was not started with TaskGroup.start()"
  312. ) from None
  313. async def wait(self) -> None:
  314. """
  315. Wait for the task to finish.
  316. This method will return as soon as the task has finished, no matter how it
  317. happened.
  318. """
  319. await self._finished_event.wait()
  320. def __await__(self) -> Generator[Any, Any, T_co]:
  321. yield from self._finished_event.wait().__await__()
  322. return self.return_value
  323. def __repr__(self) -> str:
  324. return (
  325. f"<{self.__class__.__name__} {self.status.name.lower()} "
  326. f"name={self._name!r} coro={self._coro!r}>"
  327. )