from __future__ import annotations import asyncio import time class TokenBucketRateLimiter: def __init__(self, rate: int, per: float = 60.0): self._rate = rate self._per = per self._tokens = float(rate) self._last_refill = time.monotonic() self._lock = asyncio.Lock() async def acquire(self) -> bool: async with self._lock: return self._try_acquire() def acquire_sync(self) -> bool: return self._try_acquire_sync() def _try_acquire(self) -> bool: now = time.monotonic() elapsed = now - self._last_refill self._tokens = min(self._rate, self._tokens + elapsed * (self._rate / self._per)) self._last_refill = now if self._tokens >= 1.0: self._tokens -= 1.0 return True return False def _try_acquire_sync(self) -> bool: now = time.monotonic() elapsed = now - self._last_refill self._tokens = min(self._rate, self._tokens + elapsed * (self._rate / self._per)) self._last_refill = now if self._tokens >= 1.0: self._tokens -= 1.0 return True return False @property def available_tokens(self) -> float: return self._tokens @property def rate(self) -> int: return self._rate @property def period(self) -> float: return self._per