from __future__ import annotations
import asyncio
import contextvars
import enum
import warnings
from collections.abc import Callable
from contextvars import Context
from typing import (
Any,
TypeAlias,
TypedDict,
TypeVar,
)
from .cancel import cancel_and_wait
from .types import CoroutineLike
__all__ = (
"ErrorArg",
"ErrorCallback",
"TaskContext",
)
class LoopExceptionHandler(enum.Enum):
TOKEN = 0
[docs]
class ErrorArg(TypedDict):
# Intentionally designed as a typed dict to match with
# the signature of asyncio stdlib's error callback handler.
message: str
exception: BaseException
task: asyncio.Task[Any]
ErrorCallback: TypeAlias = Callable[[ErrorArg], None]
T = TypeVar("T")
[docs]
class TaskContext:
"""
TaskContext keeps the references to the child tasks during its lifetime,
so that they can be terminated safely when shutdown is explicitly requested.
This is the loosest form of child task managers among TaskScope, TaskGroup, and
Supervisor, as it does not enforce structured concurrency but just provides a
reference set to child tasks.
You may replace existing patterns using :class:`weakref.WeakSet` to keep track
of child tasks for a long-running server application with TaskContext.
If ``exception_handler`` is not set (the default behavior), it will run
:meth:`loop.call_exception_handler() <asyncio.loop.call_exception_handler>`
with the context argument consisting of the ``message``, ``task`` (the child task
that raised the exception), and ``exception`` (the exception object) fields.
If it is set *None*, it will silently ignore the exception.
If it is set as a callable function, it will invoke the specified callback
function using the context argument same to that used when calling
:meth:`loop.call_exception_handler() <asyncio.loop.call_exception_handler>`.
If you provide ``context``, it will be passed to :meth:`create_task()` by default.
.. note::
Unlike :class:`~aiotools.taskscope.TaskScope`, TaskContext does not update
the callgraph introduced in Python 3.14 as it does not implement completion-based
waiting semantics but just grouped cancellation when explicitly requested.
.. versionadded:: 2.0
"""
# We need to keep track of strong references as the event loop keeps weak references to
# fire-and-forget tasks until Python 3.13. In Python 3.14, it will be changed as
# per-thread doubly-linked lists to support free-threaded (nogil) setups to avoid
# lock contention in the weakref subsystem.
_loop: asyncio.AbstractEventLoop
_tasks: set[asyncio.Task[Any]]
_host_task: asyncio.Task[Any] | None
_exception_handler: ErrorCallback | LoopExceptionHandler | None
_default_context: contextvars.Context | None
_entered: bool
_exited: bool
_aborting: bool
def __init__(
self,
*,
exception_handler: ErrorCallback
| LoopExceptionHandler
| None = LoopExceptionHandler.TOKEN,
context: contextvars.Context | None = None,
) -> None:
self._loop = asyncio.get_running_loop()
self._tasks = set()
self._host_task = None
self._exception_handler = exception_handler
self._default_context = context
# status flags
self._entered = False
self._exited = False
self._aborting = False
def __del__(self) -> None:
loc: str = "<loc>" # TODO: implement
if self._entered and not self._exited:
warnings.warn(
f"TaskContext initialized at {loc} is not properly "
"terminated until it is garbage-collected.",
category=ResourceWarning,
)
def __repr__(self) -> str:
info = [""]
if self._tasks:
info.append(f"tasks={len(self._tasks)}")
if self._aborting:
info.append("cancelling")
elif self._entered:
info.append("entered")
info_str = " ".join(info)
return f"<{type(self).__name__} {info_str}>"
[docs]
async def aclose(self) -> None:
"""
Triggers cancellation and waits for completion.
"""
self._aborting = True
try:
# NOTE: unhandled exceptions are captured by our own task-done handler.
await cancel_and_wait(self._tasks)
finally:
self._exited = True
[docs]
def create_task(
self,
coro: CoroutineLike[T],
*,
name: str | None = None,
context: Context | None = None,
**kwargs: Any,
) -> asyncio.Task[T]:
"""
Create a new task in this scope and return it.
Similar to :func:`asyncio.create_task()`.
"""
if not self._entered:
self._entered = True
host_task = asyncio.current_task()
assert host_task is not None
if self._host_task is not None and host_task is not self._host_task:
raise RuntimeError(
"{type(self).__name__} must be used within a single parent task."
)
self._host_task = host_task
return self._create_task(coro, name=name, context=context, **kwargs)
def _create_task(
self,
coro: CoroutineLike[T],
*,
name: str | None = None,
context: Context | None = None,
**kwargs: Any,
) -> asyncio.Task[T]:
if self._aborting:
raise RuntimeError(f"{type(self).__name__} {self!r} is shutting down")
task: asyncio.Task[T] = self._loop.create_task(
coro,
name=name,
context=self._default_context if context is None else context,
**kwargs,
)
# optimization: Immediately call the done callback if the task is
# already done (e.g. if the coro was able to complete eagerly),
# and skip scheduling a done callback
if task.done():
self._on_task_done(task)
else:
self._tasks.add(task)
task.add_done_callback(self._on_task_done)
return task
def _on_task_done(self, task: asyncio.Task[Any]) -> None:
self._tasks.discard(task)
if task.cancelled():
return
exc = task.exception()
if exc is None:
return
self._handle_task_exception(task)
def _handle_task_exception(self, task: asyncio.Task[Any]) -> None:
exc = task.exception()
assert exc is not None
match self._exception_handler:
case None:
pass # deliberately set to ignore errors
case func if callable(func):
func({
"message": (
f"Task {task!r} has errored inside the parent "
f"task {self._host_task}"
),
"exception": exc,
"task": task,
})
case LoopExceptionHandler():
self._loop.call_exception_handler({
"message": (
f"Task {task!r} has errored inside the parent "
f"task {self._host_task}"
),
"exception": exc,
"task": task,
})
case _:
raise RuntimeError(
f"Invalid error handler: {self._exception_handler!r}"
)