42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
"""Sliding-window rate limiter.
|
|
|
|
Public API:
|
|
SlidingWindow(limit, window) with .allow(key, now=None) -> bool
|
|
Counts events per key within a rolling window of `window` seconds.
|
|
Thread-safe? No (single-threaded). Deterministic given `now`.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from collections import defaultdict, deque
|
|
from typing import Deque, Dict, Optional
|
|
|
|
|
|
class SlidingWindow:
|
|
def __init__(self, limit: int, window: float):
|
|
if limit < 1:
|
|
raise ValueError("limit must be >= 1")
|
|
if window <= 0:
|
|
raise ValueError("window must be > 0")
|
|
self.limit = limit
|
|
self.window = window
|
|
self._events: Dict[object, Deque[float]] = defaultdict(deque)
|
|
|
|
def allow(self, key, now: Optional[float] = None) -> bool:
|
|
if now is None:
|
|
import time
|
|
now = time.monotonic()
|
|
q = self._events[key]
|
|
cutoff = now - self.window
|
|
while q and q[0] <= cutoff:
|
|
q.popleft()
|
|
if len(q) < self.limit:
|
|
q.append(now)
|
|
return True
|
|
return False
|
|
|
|
def reset(self, key=None):
|
|
if key is None:
|
|
self._events.clear()
|
|
else:
|
|
self._events.pop(key, None)
|