_tasks.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223
  1. from __future__ import annotations
  2. import sys
  3. from abc import ABCMeta, abstractmethod
  4. from collections.abc import Callable, Coroutine
  5. from contextvars import Context
  6. from types import TracebackType
  7. from typing import TYPE_CHECKING, Any, Literal, Protocol, final, overload
  8. if sys.version_info >= (3, 13):
  9. from typing import TypeVar
  10. else:
  11. from typing_extensions import TypeVar
  12. if sys.version_info >= (3, 11):
  13. from typing import TypeVarTuple, Unpack
  14. else:
  15. from typing_extensions import TypeVarTuple, Unpack
  16. if TYPE_CHECKING:
  17. from .._core._tasks import CancelScope, TaskHandle
  18. T_co = TypeVar("T_co", covariant=True)
  19. T_contra = TypeVar("T_contra", contravariant=True, default=None)
  20. PosArgsT = TypeVarTuple("PosArgsT")
  21. def get_coro_name(coro: Coroutine[Any, Any, object], override: object = None) -> str:
  22. if override is not None:
  23. return str(override)
  24. try:
  25. cr_frame = coro.cr_frame # type: ignore[attr-defined]
  26. module = cr_frame.f_globals["__name__"]
  27. except (AttributeError, KeyError):
  28. module = None
  29. qualname = getattr(coro, "__qualname__", None)
  30. return ".".join([x for x in (module, qualname) if x])
  31. def get_callable_name(func: Callable, override: object = None) -> str:
  32. if override is not None:
  33. return str(override)
  34. module = getattr(func, "__module__", None)
  35. qualname = getattr(func, "__qualname__", None)
  36. return ".".join([x for x in (module, qualname) if x])
  37. def call_for_coroutine(
  38. func: Callable[[Unpack[PosArgsT]], Coroutine[Any, Any, T_co]],
  39. args: tuple[Unpack[PosArgsT]],
  40. **kwargs: Any,
  41. ) -> Coroutine[Any, Any, T_co]:
  42. """
  43. Call the given function with the given positional and keyword arguments.
  44. :return: the resulting coroutine
  45. :raises TypeError: if the return value was not a coroutine object
  46. """
  47. coro = func(*args, **kwargs)
  48. if not isinstance(coro, Coroutine):
  49. prefix = f"{func.__module__}." if hasattr(func, "__module__") else ""
  50. raise TypeError(
  51. f"Expected {prefix}{func.__qualname__}() to return a coroutine, but "
  52. f"the return value ({coro!r}) is not a coroutine object"
  53. )
  54. return coro
  55. class TaskStatus(Protocol[T_contra]):
  56. @overload
  57. def started(self: TaskStatus[None]) -> None: ...
  58. @overload
  59. def started(self, value: T_contra) -> None: ...
  60. def started(self, value: T_contra | None = None) -> None:
  61. """
  62. Signal that the task has started.
  63. :param value: object passed back to the starter of the task
  64. """
  65. class TaskGroup(metaclass=ABCMeta):
  66. """
  67. Groups several asynchronous tasks together.
  68. :ivar cancel_scope: the cancel scope inherited by all child tasks
  69. :vartype cancel_scope: CancelScope
  70. .. note:: On asyncio, support for eager task factories is considered to be
  71. **experimental**. In particular, they don't follow the usual semantics of new
  72. tasks being scheduled on the next iteration of the event loop, and may thus
  73. cause unexpected behavior in code that wasn't written with such semantics in
  74. mind.
  75. """
  76. cancel_scope: CancelScope
  77. def cancel(self, reason: str | None = None) -> None:
  78. """
  79. Cancel this task group's cancel scope immediately.
  80. This is a shortcut for calling ``.cancel_scope.cancel()`` on the task group.
  81. :param reason: a message describing the reason for the cancellation
  82. .. versionadded:: 4.14.0
  83. """
  84. self.cancel_scope.cancel(reason)
  85. @abstractmethod
  86. def create_task(
  87. self,
  88. coro: Coroutine[Any, Any, T_co],
  89. *,
  90. name: object = None,
  91. context: Context | None = None,
  92. ) -> TaskHandle[T_co]:
  93. """
  94. Create a new task from a coroutine object and schedule it to run.
  95. :param coro: a coroutine object
  96. :param name: optional name to give the task
  97. :param context: optional context to run the task in
  98. :return: a task handle
  99. .. versionadded:: 4.14.0
  100. """
  101. @final
  102. def start_soon(
  103. self,
  104. func: Callable[[Unpack[PosArgsT]], Coroutine[Any, Any, T_co]],
  105. *args: Unpack[PosArgsT],
  106. name: object = None,
  107. ) -> TaskHandle[T_co]:
  108. """
  109. Start a new task in this task group.
  110. :param func: a coroutine function
  111. :param args: positional arguments to call the function with
  112. :param name: name of the task, for the purposes of introspection and debugging
  113. :return: a task handle
  114. .. versionadded:: 3.0
  115. .. versionchanged:: 4.14.0
  116. This method now returns a task handle.
  117. """
  118. final_name = get_callable_name(func, name)
  119. return self.create_task(call_for_coroutine(func, args), name=final_name)
  120. @overload
  121. async def start(
  122. self,
  123. func: Callable[..., Coroutine[Any, Any, T_co]],
  124. *args: object,
  125. name: object = None,
  126. return_handle: Literal[False] = ...,
  127. ) -> Any: ...
  128. @overload
  129. async def start(
  130. self,
  131. func: Callable[..., Coroutine[Any, Any, T_co]],
  132. *args: object,
  133. name: object = None,
  134. return_handle: Literal[True],
  135. ) -> TaskHandle[T_co, Any]: ...
  136. @abstractmethod
  137. async def start(
  138. self,
  139. func: Callable[..., Coroutine[Any, Any, T_co]],
  140. *args: object,
  141. name: object = None,
  142. return_handle: Literal[False, True] = False,
  143. ) -> Any:
  144. """
  145. Start a new task and wait until it signals for readiness.
  146. The target callable must accept a keyword argument ``task_status`` (of type
  147. :class:`TaskStatus`). Awaiting on this method will return whatever was passed to
  148. ``task_status.started()`` (``None`` by default).
  149. .. note:: The :class:`TaskStatus` class is generic, and the type argument should
  150. indicate the type of the value that will be passed to
  151. ``task_status.started()``.
  152. :param func: a coroutine function that accepts the ``task_status`` keyword
  153. argument
  154. :param args: positional arguments to call the function with
  155. :param name: an optional name for the task, for introspection and debugging
  156. :param return_handle: if ``True``, return a :class:`TaskHandle` which also
  157. contains the start value in ``start_value``
  158. :return: the value passed to ``task_status.started()``
  159. :raises RuntimeError: if the task finishes without calling
  160. ``task_status.started()``
  161. .. seealso:: :ref:`start_initialize`
  162. .. versionadded:: 3.0
  163. """
  164. @abstractmethod
  165. async def __aenter__(self) -> TaskGroup:
  166. """Enter the task group context and allow starting new tasks."""
  167. @abstractmethod
  168. async def __aexit__(
  169. self,
  170. exc_type: type[BaseException] | None,
  171. exc_val: BaseException | None,
  172. exc_tb: TracebackType | None,
  173. ) -> bool:
  174. """Exit the task group context waiting for all tasks to finish."""