Source code for uprate.ratelimit

from __future__ import annotations

from operator import attrgetter
from time import time as unix
from typing import Generic, Optional, TypeVar, Union

from import BaseStore, MemoryStore

from .errors import RateLimitError
from .rate import Rate, RateGroup

__all__ = (

H = TypeVar("H")
"""A TypeVar"""

[docs]class RateLimit(Generic[H]): """Enforces multiple rates per provided keys. This is a low-level component. Parameters ---------- rate: :class:`~uprate.rate.Rate`, :class:`~uprate.rate.RateGroup`, (``Rate | RateGroup``) The rate(s) to enforce on keys in this ratelimits. store: :class:``, (``BaseStore | None``) The store to use for this RateLimit, If None then a :class:`` is used. By default, :data:`None`. Raises ------ :exc:`TypeError` ``rate`` parameter provided to the constructor is of invalid type. Attributes ---------- rates : tuple[:class:`~uprate.rate.Rate`] A tuple of :class:`~uprate.rate.Rate` sorted in ascending order by the rate's time period. store : :class:`` The store in use for this RateLimit. """ rates: tuple[Rate, ...] store: BaseStore[H] def __init__(self, rate: Union[Rate, RateGroup], store: Optional[BaseStore[H]] = None) -> None: if isinstance(rate, Rate): self.rates = (rate,) elif isinstance(rate, RateGroup): rate._data.sort(key=attrgetter("period")) self.rates = tuple(rate._data) else: raise TypeError(f"Expected instance of uprate.rate.Rate or uprate.rate.RateGroup Instead got {type(rate)}") if store is None: = MemoryStore() elif isinstance(store, BaseStore): = store else: raise TypeError("Expected a type deriving from instead got " + str(type(store)))
[docs] async def acquire(self, key: H) -> None: """Try to acquire a usage token for given token. Raise an error if token can't be acquired. Parameters ---------- key : :data:`.H` The key to acquire a usage token for. Raises ------ :exc:`~uprate.errors.RateLimitError` Cannot acquire usage token due to exhaustion. """ res, retry, rate = await if not res: # cast is ugly, overloads don't work (parameters don't change) # TODO: raise RateLimitError(retry_at=retry + unix(), rate=rate) # type: ignore[arg-type]
[docs] async def reset(self, key: H = None) -> None: """Reset the given key. Parameters ---------- key : :data:`.H`, :data:`None`, (``H | None``) The key to reset ratelimit for. If :data:`None`, then resets all ratelimits, by default :data:`None`. """ if not key: await else: await