|
- import enum
- import sys
- from dataclasses import dataclass
- from typing import Any, Dict, Generic, Set, TypeVar, Union, overload
- from weakref import WeakKeyDictionary
-
- from ._core._eventloop import get_asynclib
-
- if sys.version_info >= (3, 8):
- from typing import Literal
- else:
- from typing_extensions import Literal
-
- 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_asynclib().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_asynclib().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_asynclib().cancel_shielded_checkpoint()
-
-
- def current_token() -> object:
- """Return a backend specific token object that can be used to get back to the event loop."""
- return get_asynclib().current_token()
-
-
- _run_vars = WeakKeyDictionary() # type: WeakKeyDictionary[Any, Dict[str, Any]]
- _token_wrappers: Dict[Any, '_TokenWrapper'] = {}
-
-
- @dataclass(frozen=True)
- class _TokenWrapper:
- __slots__ = '_token', '__weakref__'
- _token: object
-
-
- class _NoValueSet(enum.Enum):
- NO_VALUE_SET = enum.auto()
-
-
- class RunvarToken(Generic[T]):
- __slots__ = '_var', '_value', '_redeemed'
-
- def __init__(self, var: 'RunVar[T]', value: Union[T, Literal[_NoValueSet.NO_VALUE_SET]]):
- self._var = var
- self._value: Union[T, Literal[_NoValueSet.NO_VALUE_SET]] = value
- self._redeemed = False
-
-
- class RunVar(Generic[T]):
- """Like a :class:`~contextvars.ContextVar`, expect scoped to the running event loop."""
- __slots__ = '_name', '_default'
-
- NO_VALUE_SET: Literal[_NoValueSet.NO_VALUE_SET] = _NoValueSet.NO_VALUE_SET
-
- _token_wrappers: Set[_TokenWrapper] = set()
-
- def __init__(self, name: str,
- default: Union[T, Literal[_NoValueSet.NO_VALUE_SET]] = NO_VALUE_SET):
- self._name = name
- self._default = default
-
- @property
- def _current_vars(self) -> Dict[str, T]:
- token = current_token()
- while True:
- try:
- return _run_vars[token]
- except TypeError:
- # Happens when token isn't weak referable (TrioToken).
- # This workaround does mean that some memory will leak on Trio until the problem
- # is fixed on their end.
- token = _TokenWrapper(token)
- self._token_wrappers.add(token)
- except KeyError:
- run_vars = _run_vars[token] = {}
- return run_vars
-
- @overload
- def get(self, default: D) -> Union[T, D]: ...
-
- @overload
- def get(self) -> T: ...
-
- def get(
- self, default: Union[D, Literal[_NoValueSet.NO_VALUE_SET]] = NO_VALUE_SET
- ) -> Union[T, D]:
- try:
- return self._current_vars[self._name]
- 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]:
- current_vars = self._current_vars
- token = RunvarToken(self, current_vars.get(self._name, RunVar.NO_VALUE_SET))
- current_vars[self._name] = value
- return token
-
- def reset(self, token: RunvarToken[T]) -> None:
- 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._name]
- except KeyError:
- pass
- else:
- self._current_vars[self._name] = token._value
-
- token._redeemed = True
-
- def __repr__(self) -> str:
- return f'<RunVar name={self._name!r}>'
|