Token Bucket Rate Limiter in Python

import threading
import time
from collections.abc import Callable
from dataclasses import dataclass, field


@dataclass
class TokenBucket:
    capacity: float
    refill_per_second: float
    clock: Callable[[], float] = time.monotonic
    _tokens: float = field(init=False)
    _updated: float = field(init=False)
    _lock: threading.Lock = field(
        default_factory=threading.Lock,
        init=False,
        repr=False,
    )

    def __post_init__(self) -> None:
        if self.capacity <= 0:
            raise ValueError("capacity must be positive")
        if self.refill_per_second <= 0:
            raise ValueError("refill rate must be positive")
        self._tokens = self.capacity
        self._updated = self.clock()

    def _refill(self) -> None:
        now = self.clock()
        elapsed = max(0.0, now - self._updated)
        self._tokens = min(
            self.capacity,
            self._tokens + elapsed * self.refill_per_second,
        )
        self._updated = now

    def allow(self, cost: float = 1.0) -> tuple[bool, float]:
        if cost <= 0 or cost > self.capacity:
            raise ValueError("cost must be within (0, capacity]")
        with self._lock:
            self._refill()
            if self._tokens >= cost:
                self._tokens -= cost
                return True, 0.0
            missing = cost - self._tokens
            return False, missing / self.refill_per_second


class FakeClock:
    def __init__(self) -> None:
        self.current = 1_000.0

    def __call__(self) -> float:
        return self.current

    def advance(self, seconds: float) -> None:
        self.current += seconds


def main() -> None:
    clock = FakeClock()
    bucket = TokenBucket(capacity=3, refill_per_second=2, clock=clock)
    for request in range(1, 5):
        allowed, wait = bucket.allow()
        print(request, allowed, round(wait, 2))
    clock.advance(0.5)
    print("after 0.5s:", bucket.allow())


if __name__ == "__main__":
    main()