_tasks.py 6.5 KB

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