File size: 6,280 Bytes
ea35292 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 | from __future__ import annotations
__all__ = (
"EventLoopToken",
"RunvarToken",
"RunVar",
"checkpoint",
"checkpoint_if_cancelled",
"cancel_shielded_checkpoint",
"current_token",
)
import enum
from dataclasses import dataclass
from types import TracebackType
from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, final, overload
from weakref import WeakKeyDictionary
from ._core._eventloop import get_async_backend
if TYPE_CHECKING:
from .abc import AsyncBackend
T = TypeVar("T")
D = TypeVar("D")
async def checkpoint() -> None:
"""
Check for cancellation and allow the scheduler to switch to another task.
Equivalent to (but more efficient than)::
await checkpoint_if_cancelled()
await cancel_shielded_checkpoint()
.. versionadded:: 3.0
"""
await get_async_backend().checkpoint()
async def checkpoint_if_cancelled() -> None:
"""
Enter a checkpoint if the enclosing cancel scope has been cancelled.
This does not allow the scheduler to switch to a different task.
.. versionadded:: 3.0
"""
await get_async_backend().checkpoint_if_cancelled()
async def cancel_shielded_checkpoint() -> None:
"""
Allow the scheduler to switch to another task but without checking for cancellation.
Equivalent to (but potentially more efficient than)::
with CancelScope(shield=True):
await checkpoint()
.. versionadded:: 3.0
"""
await get_async_backend().cancel_shielded_checkpoint()
@final
@dataclass(frozen=True, repr=False)
class EventLoopToken:
"""
An opaque object that holds a reference to an event loop.
.. versionadded:: 4.11.0
"""
backend_class: type[AsyncBackend]
native_token: object
def current_token() -> EventLoopToken:
"""
Return a token object that can be used to call code in the current event loop from
another thread.
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
current thread
.. versionadded:: 4.11.0
"""
backend_class = get_async_backend()
raw_token = backend_class.current_token()
return EventLoopToken(backend_class, raw_token)
_run_vars: WeakKeyDictionary[object, dict[RunVar[Any], Any]] = WeakKeyDictionary()
class _NoValueSet(enum.Enum):
NO_VALUE_SET = enum.auto()
class RunvarToken(Generic[T]):
"""
A token that can be used to restore a :class:`RunVar` to its previous value.
Returned by :meth:`RunVar.set`. Can be used as a context manager to automatically
reset the variable on exit, or passed directly to :meth:`RunVar.reset`.
"""
__slots__ = "_var", "_value", "_redeemed"
def __init__(self, var: RunVar[T], value: T | Literal[_NoValueSet.NO_VALUE_SET]):
self._var = var
self._value: T | Literal[_NoValueSet.NO_VALUE_SET] = value
self._redeemed = False
def __enter__(self) -> RunvarToken[T]:
return self
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> None:
self._var.reset(self)
class RunVar(Generic[T]):
"""
Like a :class:`~contextvars.ContextVar`, except scoped to the running event loop.
Can be used as a context manager, Just like :class:`~contextvars.ContextVar`, that
will reset the variable to its previous value when the context block is exited.
"""
__slots__ = "_name", "_default"
NO_VALUE_SET: Literal[_NoValueSet.NO_VALUE_SET] = _NoValueSet.NO_VALUE_SET
def __init__(
self, name: str, default: T | Literal[_NoValueSet.NO_VALUE_SET] = NO_VALUE_SET
):
self._name = name
self._default = default
@property
def _current_vars(self) -> dict[RunVar[T], T]:
native_token = current_token().native_token
try:
return _run_vars[native_token]
except KeyError:
run_vars = _run_vars[native_token] = {}
return run_vars
@overload
def get(self, default: D) -> T | D: ...
@overload
def get(self) -> T: ...
def get(
self, default: D | Literal[_NoValueSet.NO_VALUE_SET] = NO_VALUE_SET
) -> T | D:
"""
Return the current value of this run variable.
:param default: a fallback value to return if no value has been set
:return: the current value, the provided default, or the variable's own default
:raises LookupError: if no value is set and no default is available
"""
try:
return self._current_vars[self]
except KeyError:
if default is not RunVar.NO_VALUE_SET:
return default
elif self._default is not RunVar.NO_VALUE_SET:
return self._default
raise LookupError(
f'Run variable "{self._name}" has no value and no default set'
)
def set(self, value: T) -> RunvarToken[T]:
"""
Set the value of this run variable for the current event loop.
:param value: the new value
:return: a token that can be used to restore the previous value
"""
current_vars = self._current_vars
token = RunvarToken(self, current_vars.get(self, RunVar.NO_VALUE_SET))
current_vars[self] = value
return token
def reset(self, token: RunvarToken[T]) -> None:
"""
Restore this run variable to the value it held before the matching :meth:`set`.
:param token: the token returned by :meth:`set`
:raises ValueError: if the token belongs to a different :class:`RunVar` or the token
has already been used
"""
if token._var is not self:
raise ValueError("This token does not belong to this RunVar")
if token._redeemed:
raise ValueError("This token has already been used")
if token._value is _NoValueSet.NO_VALUE_SET:
try:
del self._current_vars[self]
except KeyError:
pass
else:
self._current_vars[self] = token._value
token._redeemed = True
def __repr__(self) -> str:
return f"<RunVar name={self._name!r}>"
|