| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772 |
- from __future__ import annotations
- import math
- from collections import deque
- from collections.abc import Callable
- from dataclasses import dataclass
- from types import TracebackType
- from typing import TypeVar
- from ..lowlevel import checkpoint_if_cancelled
- from ._eventloop import get_async_backend
- from ._exceptions import BusyResourceError, NoEventLoopError
- from ._tasks import CancelScope
- from ._testing import TaskInfo, get_current_task
- T = TypeVar("T")
- @dataclass(frozen=True)
- class EventStatistics:
- """
- :ivar int tasks_waiting: number of tasks waiting on :meth:`~.Event.wait`
- """
- tasks_waiting: int
- @dataclass(frozen=True)
- class CapacityLimiterStatistics:
- """
- :ivar int borrowed_tokens: number of tokens currently borrowed by tasks
- :ivar float total_tokens: total number of available tokens
- :ivar tuple borrowers: tasks or other objects currently holding tokens borrowed from
- this limiter
- :ivar int tasks_waiting: number of tasks waiting on
- :meth:`~.CapacityLimiter.acquire` or
- :meth:`~.CapacityLimiter.acquire_on_behalf_of`
- """
- borrowed_tokens: int
- total_tokens: float
- borrowers: tuple[object, ...]
- tasks_waiting: int
- @dataclass(frozen=True)
- class LockStatistics:
- """
- :ivar bool locked: flag indicating if this lock is locked or not
- :ivar ~anyio.TaskInfo owner: task currently holding the lock (or ``None`` if the
- lock is not held by any task)
- :ivar int tasks_waiting: number of tasks waiting on :meth:`~.Lock.acquire`
- """
- locked: bool
- owner: TaskInfo | None
- tasks_waiting: int
- @dataclass(frozen=True)
- class ConditionStatistics:
- """
- :ivar int tasks_waiting: number of tasks blocked on :meth:`~.Condition.wait`
- :ivar ~anyio.LockStatistics lock_statistics: statistics of the underlying
- :class:`~.Lock`
- """
- tasks_waiting: int
- lock_statistics: LockStatistics
- @dataclass(frozen=True)
- class SemaphoreStatistics:
- """
- :ivar int tasks_waiting: number of tasks waiting on :meth:`~.Semaphore.acquire`
- """
- tasks_waiting: int
- class Event:
- __slots__ = ("__weakref__",)
- def __new__(cls) -> Event:
- try:
- return get_async_backend().create_event()
- except NoEventLoopError:
- return EventAdapter()
- def set(self) -> None:
- """Set the flag, notifying all listeners."""
- raise NotImplementedError
- def is_set(self) -> bool:
- """Return ``True`` if the flag is set, ``False`` if not."""
- raise NotImplementedError
- async def wait(self) -> None:
- """
- Wait until the flag has been set.
- If the flag has already been set when this method is called, it returns
- immediately.
- """
- raise NotImplementedError
- def statistics(self) -> EventStatistics:
- """Return statistics about the current state of this event."""
- raise NotImplementedError
- class EventAdapter(Event):
- __slots__ = "_internal_event", "_is_set"
- def __new__(cls) -> EventAdapter:
- return object.__new__(cls)
- def __init__(self) -> None:
- self._internal_event: Event | None = None
- self._is_set = False
- @property
- def _event(self) -> Event:
- if self._internal_event is None:
- self._internal_event = get_async_backend().create_event()
- if self._is_set:
- self._internal_event.set()
- return self._internal_event
- def set(self) -> None:
- if self._internal_event is None:
- self._is_set = True
- else:
- self._event.set()
- def is_set(self) -> bool:
- if self._internal_event is None:
- return self._is_set
- return self._internal_event.is_set()
- async def wait(self) -> None:
- await self._event.wait()
- def statistics(self) -> EventStatistics:
- if self._internal_event is None:
- return EventStatistics(tasks_waiting=0)
- return self._internal_event.statistics()
- class Lock:
- __slots__ = ("__weakref__",)
- def __new__(cls, *, fast_acquire: bool = False) -> Lock:
- try:
- return get_async_backend().create_lock(fast_acquire=fast_acquire)
- except NoEventLoopError:
- return LockAdapter(fast_acquire=fast_acquire)
- async def __aenter__(self) -> None:
- await self.acquire()
- async def __aexit__(
- self,
- exc_type: type[BaseException] | None,
- exc_val: BaseException | None,
- exc_tb: TracebackType | None,
- ) -> None:
- self.release()
- async def acquire(self) -> None:
- """Acquire the lock."""
- raise NotImplementedError
- def acquire_nowait(self) -> None:
- """
- Acquire the lock, without blocking.
- :raises ~anyio.WouldBlock: if the operation would block
- """
- raise NotImplementedError
- def release(self) -> None:
- """Release the lock."""
- raise NotImplementedError
- def locked(self) -> bool:
- """Return True if the lock is currently held."""
- raise NotImplementedError
- def statistics(self) -> LockStatistics:
- """
- Return statistics about the current state of this lock.
- .. versionadded:: 3.0
- """
- raise NotImplementedError
- class LockAdapter(Lock):
- __slots__ = "_internal_lock", "_fast_acquire"
- def __new__(cls, *, fast_acquire: bool = False) -> LockAdapter:
- return object.__new__(cls)
- def __init__(self, *, fast_acquire: bool = False):
- self._internal_lock: Lock | None = None
- self._fast_acquire = fast_acquire
- @property
- def _lock(self) -> Lock:
- if self._internal_lock is None:
- self._internal_lock = get_async_backend().create_lock(
- fast_acquire=self._fast_acquire
- )
- return self._internal_lock
- async def __aenter__(self) -> None:
- await self._lock.acquire()
- async def __aexit__(
- self,
- exc_type: type[BaseException] | None,
- exc_val: BaseException | None,
- exc_tb: TracebackType | None,
- ) -> None:
- if self._internal_lock is not None:
- self._internal_lock.release()
- async def acquire(self) -> None:
- """Acquire the lock."""
- await self._lock.acquire()
- def acquire_nowait(self) -> None:
- """
- Acquire the lock, without blocking.
- :raises ~anyio.WouldBlock: if the operation would block
- """
- self._lock.acquire_nowait()
- def release(self) -> None:
- """Release the lock."""
- self._lock.release()
- def locked(self) -> bool:
- """Return True if the lock is currently held."""
- return self._lock.locked()
- def statistics(self) -> LockStatistics:
- """
- Return statistics about the current state of this lock.
- .. versionadded:: 3.0
- """
- if self._internal_lock is None:
- return LockStatistics(False, None, 0)
- return self._internal_lock.statistics()
- class Condition:
- __slots__ = "__weakref__", "_owner_task", "_lock", "_waiters"
- def __init__(self, lock: Lock | None = None):
- self._owner_task: TaskInfo | None = None
- self._lock = lock or Lock()
- self._waiters: deque[Event] = deque()
- async def __aenter__(self) -> None:
- await self.acquire()
- async def __aexit__(
- self,
- exc_type: type[BaseException] | None,
- exc_val: BaseException | None,
- exc_tb: TracebackType | None,
- ) -> None:
- self.release()
- def _check_acquired(self) -> None:
- if self._owner_task != get_current_task():
- raise RuntimeError("The current task is not holding the underlying lock")
- async def acquire(self) -> None:
- """Acquire the underlying lock."""
- await self._lock.acquire()
- self._owner_task = get_current_task()
- def acquire_nowait(self) -> None:
- """
- Acquire the underlying lock, without blocking.
- :raises ~anyio.WouldBlock: if the operation would block
- """
- self._lock.acquire_nowait()
- self._owner_task = get_current_task()
- def release(self) -> None:
- """Release the underlying lock."""
- self._lock.release()
- def locked(self) -> bool:
- """Return True if the lock is set."""
- return self._lock.locked()
- def notify(self, n: int = 1) -> None:
- """Notify exactly n listeners."""
- self._check_acquired()
- for _ in range(n):
- try:
- event = self._waiters.popleft()
- except IndexError:
- break
- event.set()
- def notify_all(self) -> None:
- """Notify all the listeners."""
- self._check_acquired()
- for event in self._waiters:
- event.set()
- self._waiters.clear()
- async def wait(self) -> None:
- """Wait for a notification."""
- await checkpoint_if_cancelled()
- self._check_acquired()
- event = Event()
- self._waiters.append(event)
- self.release()
- try:
- await event.wait()
- except BaseException:
- if not event.is_set():
- self._waiters.remove(event)
- elif self._waiters:
- # This task was notified by could not act on it, so pass
- # it on to the next task
- self._waiters.popleft().set()
- raise
- finally:
- with CancelScope(shield=True):
- await self.acquire()
- async def wait_for(self, predicate: Callable[[], T]) -> T:
- """
- Wait until a predicate becomes true.
- :param predicate: a callable that returns a truthy value when the condition is
- met
- :return: the result of the predicate
- .. versionadded:: 4.11.0
- """
- while not (result := predicate()):
- await self.wait()
- return result
- def statistics(self) -> ConditionStatistics:
- """
- Return statistics about the current state of this condition.
- .. versionadded:: 3.0
- """
- return ConditionStatistics(len(self._waiters), self._lock.statistics())
- class Semaphore:
- __slots__ = "__weakref__", "_fast_acquire"
- def __new__(
- cls,
- initial_value: int,
- *,
- max_value: int | None = None,
- fast_acquire: bool = False,
- ) -> Semaphore:
- try:
- return get_async_backend().create_semaphore(
- initial_value, max_value=max_value, fast_acquire=fast_acquire
- )
- except NoEventLoopError:
- return SemaphoreAdapter(initial_value, max_value=max_value)
- def __init__(
- self,
- initial_value: int,
- *,
- max_value: int | None = None,
- fast_acquire: bool = False,
- ):
- if not isinstance(initial_value, int):
- raise TypeError("initial_value must be an integer")
- if initial_value < 0:
- raise ValueError("initial_value must be >= 0")
- if max_value is not None:
- if not isinstance(max_value, int):
- raise TypeError("max_value must be an integer or None")
- if max_value < initial_value:
- raise ValueError(
- "max_value must be equal to or higher than initial_value"
- )
- self._fast_acquire = fast_acquire
- async def __aenter__(self) -> Semaphore:
- await self.acquire()
- return self
- async def __aexit__(
- self,
- exc_type: type[BaseException] | None,
- exc_val: BaseException | None,
- exc_tb: TracebackType | None,
- ) -> None:
- self.release()
- async def acquire(self) -> None:
- """Decrement the semaphore value, blocking if necessary."""
- raise NotImplementedError
- def acquire_nowait(self) -> None:
- """
- Acquire the underlying lock, without blocking.
- :raises ~anyio.WouldBlock: if the operation would block
- """
- raise NotImplementedError
- def release(self) -> None:
- """Increment the semaphore value."""
- raise NotImplementedError
- @property
- def value(self) -> int:
- """The current value of the semaphore."""
- raise NotImplementedError
- @property
- def max_value(self) -> int | None:
- """The maximum value of the semaphore."""
- raise NotImplementedError
- def statistics(self) -> SemaphoreStatistics:
- """
- Return statistics about the current state of this semaphore.
- .. versionadded:: 3.0
- """
- raise NotImplementedError
- class SemaphoreAdapter(Semaphore):
- __slots__ = "_internal_semaphore", "_initial_value", "_max_value"
- def __new__(
- cls,
- initial_value: int,
- *,
- max_value: int | None = None,
- fast_acquire: bool = False,
- ) -> SemaphoreAdapter:
- return object.__new__(cls)
- def __init__(
- self,
- initial_value: int,
- *,
- max_value: int | None = None,
- fast_acquire: bool = False,
- ) -> None:
- super().__init__(initial_value, max_value=max_value, fast_acquire=fast_acquire)
- self._internal_semaphore: Semaphore | None = None
- self._initial_value = initial_value
- self._max_value = max_value
- @property
- def _semaphore(self) -> Semaphore:
- if self._internal_semaphore is None:
- self._internal_semaphore = get_async_backend().create_semaphore(
- self._initial_value, max_value=self._max_value
- )
- return self._internal_semaphore
- async def acquire(self) -> None:
- await self._semaphore.acquire()
- def acquire_nowait(self) -> None:
- self._semaphore.acquire_nowait()
- def release(self) -> None:
- self._semaphore.release()
- @property
- def value(self) -> int:
- if self._internal_semaphore is None:
- return self._initial_value
- return self._semaphore.value
- @property
- def max_value(self) -> int | None:
- return self._max_value
- def statistics(self) -> SemaphoreStatistics:
- if self._internal_semaphore is None:
- return SemaphoreStatistics(tasks_waiting=0)
- return self._semaphore.statistics()
- class CapacityLimiter:
- __slots__ = ("__weakref__",)
- def __new__(cls, total_tokens: float) -> CapacityLimiter:
- try:
- return get_async_backend().create_capacity_limiter(total_tokens)
- except NoEventLoopError:
- return CapacityLimiterAdapter(total_tokens)
- async def __aenter__(self) -> None:
- raise NotImplementedError
- async def __aexit__(
- self,
- exc_type: type[BaseException] | None,
- exc_val: BaseException | None,
- exc_tb: TracebackType | None,
- ) -> None:
- raise NotImplementedError
- @property
- def total_tokens(self) -> float:
- """
- The total number of tokens available for borrowing.
- This is a read-write property. If the total number of tokens is increased, the
- proportionate number of tasks waiting on this limiter will be granted their
- tokens.
- .. versionchanged:: 3.0
- The property is now writable.
- .. versionchanged:: 4.12
- The value can now be set to 0.
- """
- raise NotImplementedError
- @total_tokens.setter
- def total_tokens(self, value: float) -> None:
- raise NotImplementedError
- @property
- def borrowed_tokens(self) -> int:
- """The number of tokens that have currently been borrowed."""
- raise NotImplementedError
- @property
- def available_tokens(self) -> float:
- """The number of tokens currently available to be borrowed"""
- raise NotImplementedError
- def acquire_nowait(self) -> None:
- """
- Acquire a token for the current task without waiting for one to become
- available.
- :raises ~anyio.WouldBlock: if there are no tokens available for borrowing
- """
- raise NotImplementedError
- def acquire_on_behalf_of_nowait(self, borrower: object) -> None:
- """
- Acquire a token without waiting for one to become available.
- :param borrower: the entity borrowing a token
- :raises ~anyio.WouldBlock: if there are no tokens available for borrowing
- """
- raise NotImplementedError
- async def acquire(self) -> None:
- """
- Acquire a token for the current task, waiting if necessary for one to become
- available.
- """
- raise NotImplementedError
- async def acquire_on_behalf_of(self, borrower: object) -> None:
- """
- Acquire a token, waiting if necessary for one to become available.
- :param borrower: the entity borrowing a token
- """
- raise NotImplementedError
- def release(self) -> None:
- """
- Release the token held by the current task.
- :raises RuntimeError: if the current task has not borrowed a token from this
- limiter.
- """
- raise NotImplementedError
- def release_on_behalf_of(self, borrower: object) -> None:
- """
- Release the token held by the given borrower.
- :raises RuntimeError: if the borrower has not borrowed a token from this
- limiter.
- """
- raise NotImplementedError
- def statistics(self) -> CapacityLimiterStatistics:
- """
- Return statistics about the current state of this limiter.
- .. versionadded:: 3.0
- """
- raise NotImplementedError
- class CapacityLimiterAdapter(CapacityLimiter):
- __slots__ = "_internal_limiter", "_total_tokens"
- def __new__(cls, total_tokens: float) -> CapacityLimiterAdapter:
- return object.__new__(cls)
- def __init__(self, total_tokens: float) -> None:
- self._internal_limiter: CapacityLimiter | None = None
- self.total_tokens = total_tokens
- @property
- def _limiter(self) -> CapacityLimiter:
- if self._internal_limiter is None:
- self._internal_limiter = get_async_backend().create_capacity_limiter(
- self._total_tokens
- )
- return self._internal_limiter
- async def __aenter__(self) -> None:
- await self._limiter.__aenter__()
- async def __aexit__(
- self,
- exc_type: type[BaseException] | None,
- exc_val: BaseException | None,
- exc_tb: TracebackType | None,
- ) -> None:
- return await self._limiter.__aexit__(exc_type, exc_val, exc_tb)
- @property
- def total_tokens(self) -> float:
- if self._internal_limiter is None:
- return self._total_tokens
- return self._internal_limiter.total_tokens
- @total_tokens.setter
- def total_tokens(self, value: float) -> None:
- if not isinstance(value, int) and not math.isinf(value):
- raise TypeError("total_tokens must be an int or math.inf")
- elif value < 0:
- raise ValueError("total_tokens must be >= 0")
- if self._internal_limiter is None:
- self._total_tokens = value
- return
- self._limiter.total_tokens = value
- @property
- def borrowed_tokens(self) -> int:
- if self._internal_limiter is None:
- return 0
- return self._internal_limiter.borrowed_tokens
- @property
- def available_tokens(self) -> float:
- if self._internal_limiter is None:
- return self._total_tokens
- return self._internal_limiter.available_tokens
- def acquire_nowait(self) -> None:
- self._limiter.acquire_nowait()
- def acquire_on_behalf_of_nowait(self, borrower: object) -> None:
- self._limiter.acquire_on_behalf_of_nowait(borrower)
- async def acquire(self) -> None:
- await self._limiter.acquire()
- async def acquire_on_behalf_of(self, borrower: object) -> None:
- await self._limiter.acquire_on_behalf_of(borrower)
- def release(self) -> None:
- self._limiter.release()
- def release_on_behalf_of(self, borrower: object) -> None:
- self._limiter.release_on_behalf_of(borrower)
- def statistics(self) -> CapacityLimiterStatistics:
- if self._internal_limiter is None:
- return CapacityLimiterStatistics(
- borrowed_tokens=0,
- total_tokens=self.total_tokens,
- borrowers=(),
- tasks_waiting=0,
- )
- return self._internal_limiter.statistics()
- class ResourceGuard:
- """
- A context manager for ensuring that a resource is only used by a single task at a
- time.
- Entering this context manager while the previous has not exited it yet will trigger
- :exc:`BusyResourceError`.
- :param action: the action to guard against (visible in the :exc:`BusyResourceError`
- when triggered, e.g. "Another task is already {action} this resource")
- .. versionadded:: 4.1
- """
- __slots__ = "__weakref__", "action", "_guarded"
- def __init__(self, action: str = "using"):
- self.action: str = action
- self._guarded = False
- def __enter__(self) -> None:
- if self._guarded:
- raise BusyResourceError(self.action)
- self._guarded = True
- def __exit__(
- self,
- exc_type: type[BaseException] | None,
- exc_val: BaseException | None,
- exc_tb: TracebackType | None,
- ) -> None:
- self._guarded = False
|