Token bucket rate limiter for threads

import threading
import time


class TokenBucket:
    """Allow bursts up to `capacity` while enforcing an average `rate`."""

    def __init__(self, rate: float, capacity: int) -> None:
        self.rate = rate
        self.capacity = capacity
        self.tokens = float(capacity)
        self.updated = time.monotonic()
        self.lock = threading.Lock()

    def acquire(self, amount: int = 1) -> float:
        """Block until `amount` tokens exist; return the seconds spent waiting."""
        if amount > self.capacity:
            raise ValueError("amount exceeds bucket capacity")
        waited = 0.0
        while True:
            with self.lock:
                now = time.monotonic()
                self.tokens = min(self.capacity, self.tokens + (now - self.updated) * self.rate)
                self.updated = now
                if self.tokens >= amount:
                    self.tokens -= amount
                    return waited
                missing = (amount - self.tokens) / self.rate
            time.sleep(missing)
            waited += missing


if __name__ == "__main__":
    bucket = TokenBucket(rate=2, capacity=5)
    print([round(bucket.acquire(), 2) for _ in range(6)])