from __future__ import annotations
from operator import attrgetter
from time import time as unix
from typing import Generic, Optional, TypeVar, Union
from uprate.store import BaseStore, MemoryStore
from .errors import RateLimitError
from .rate import Rate, RateGroup
__all__ = (
"RateLimit",
)
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:`~uprate.store.BaseStore`, (``BaseStore | None``)
The store to use for this RateLimit, If None then a :class:`uprate.store.MemoryStore` 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:`~uprate.store.BaseStore`
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:
self.store = MemoryStore()
elif isinstance(store, BaseStore):
self.store = store
else:
raise TypeError("Expected a type deriving from uprate.store.BaseStore instead got " + str(type(store)))
self.store.setup(self)
[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 self.store.acquire(key)
if not res:
# cast is ugly, overloads don't work (parameters don't change)
# TODO: https://www.python.org/dev/peps/pep-0647/
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 self.store.clear()
else:
await self.store.reset(key)