| 1 | import asyncio |
| 2 | import time |
| 3 | from typing import Callable, Awaitable |
| 4 | |
| 5 | |
| 6 | class RateLimiter: |
| 7 | def __init__(self, seconds: int = 60, **limits: int): |
| 8 | self.timeframe = seconds |
| 9 | self.limits = {key: value if isinstance(value, (int, float)) else 0 for key, value in (limits or {}).items()} |
| 10 | self.values = {key: [] for key in self.limits.keys()} |
| 11 | self._lock = asyncio.Lock() |
| 12 | |
| 13 | def add(self, **kwargs: int): |
| 14 | now = time.time() |
| 15 | for key, value in kwargs.items(): |
| 16 | if not key in self.values: |
| 17 | self.values[key] = [] |
| 18 | self.values[key].append((now, value)) |
| 19 | |
| 20 | async def cleanup(self): |
| 21 | async with self._lock: |
| 22 | now = time.time() |
| 23 | cutoff = now - self.timeframe |
| 24 | for key in self.values: |
| 25 | self.values[key] = [(t, v) for t, v in self.values[key] if t > cutoff] |
| 26 | |
| 27 | async def get_total(self, key: str) -> int: |
| 28 | async with self._lock: |
| 29 | if not key in self.values: |
| 30 | return 0 |
| 31 | return sum(value for _, value in self.values[key]) |
| 32 | |
| 33 | async def wait( |
| 34 | self, |
| 35 | callback: Callable[[str, str, int, int], Awaitable[bool]] | None = None, |
| 36 | ): |
| 37 | while True: |
| 38 | await self.cleanup() |
| 39 | should_wait = False |
| 40 | |
| 41 | for key, limit in self.limits.items(): |
| 42 | if limit <= 0: # Skip if no limit set |
| 43 | continue |
| 44 | |
| 45 | total = await self.get_total(key) |
| 46 | if total > limit: |
| 47 | if callback: |
| 48 | msg = f"Rate limit exceeded for {key} ({total}/{limit}), waiting..." |
| 49 | should_wait = not await callback(msg, key, total, limit) |
| 50 | else: |
| 51 | should_wait = True |
| 52 | break |
| 53 | |
| 54 | if not should_wait: |
| 55 | break |
| 56 | |
| 57 | await asyncio.sleep(1) |