TokenBucket レートリミッタクラスを実装する

import time class TokenBucket: """A token bucket rate limiter. Tokens are added at a constant `refill_rate` (tokens per second) up to `capacity`. Each call to `allow_request()` tries to consume one token. """ def __init__(self, capacity: int, refill_rate: float) -> None: """ Args: capacity: Maximum number of tokens the bucket can hold. refill_rate: Number of tokens added to the bucket per second. """ self.capacity: int = capacity self.refill_rate: float = refill_rate self.tokens: float = float(capacity) # start full self.last_refill_time: float = time.monotonic() def _refill(self) -> None: """Add tokens based on elapsed time since the last refill.""" now = time.monotonic() elapsed = now - self.last_refill_time tokens_to_add = elapsed * self.refill_rate self.tokens = min(self.capacity, self.tokens + tokens_to_add) self.last_refill_time = now def allow_request(self) -> bool: """Attempt to consume one token. Returns: True if the request is allowed (a token was available and consumed), False otherwise. """ self._refill() if self.tokens >= 1: self.tokens -= 1 return True return False # --------------------------------------------------------------------------- # Quick demo / smoke test # --------------------------------------------------------------------------- if __name__ == "__main__": bucket = TokenBucket(capacity=5, refill_rate=2) # 5 max, 2 tokens/sec # Burst: should allow 5 requests immediately for i in range(7): result = bucket.allow_request() print(f"Request {i + 1}: {'allowed' if result else 'denied'}") # Wait a bit for tokens to refill print("\nSleeping 1.5 seconds...") time.sleep(1.5) # should refill ~3 tokens for i in range(5): result = bucket.allow_request() print(f"Request {i + 1}: {'allowed' if result else 'denied'}")

import time import threading class TokenBucket: """ A simple token bucket rate limiter. Args: capacity (int): maximum number of tokens in the bucket. refill_rate (float): tokens added per second. Methods: allow_request() -> bool: returns True and consumes one token if available, otherwise False. """ def __init__(self, capacity, refill_rate): if capacity <= 0: raise ValueError("capacity must be > 0") if refill_rate < 0: raise ValueError("refill_rate must be >= 0") self.capacity = float(capacity) self.refill_rate = float(refill_rate) # current tokens may be fractional to account for partial refill self._tokens = float(capacity) self._last_time = time.monotonic() self._lock = threading.Lock() def _refill(self): """Refill tokens based on time elapsed since last refill.""" now = time.monotonic() elapsed = now - self._last_time if elapsed <= 0: return added = elapsed * self.refill_rate # cap tokens at capacity self._tokens = min(self.capacity, self._tokens + added) self._last_time = now def allow_request(self): """ Attempt to consume one token. Returns: bool: True if a token was available and consumed, False otherwise. """ with self._lock: self._refill() if self._tokens >= 1.0: self._tokens -= 1.0 return True return False # Example usage: # bucket = TokenBucket(capacity=10, refill_rate=1.5) # 1.5 tokens per second # if bucket.allow_request(): # # process request # else: # # reject or throttle request