added Token Bucket rate limiter algorithm

other/cheap_progress.py

@@ -0,0 +1,54 @@
+from collections.abc import Iterable, Iterator
+from sys import stderr
+
+
+def progress[T](items: Iterable[T], desc: str = "", total: int = 0) -> Iterator[T]:
+    """
+    A simple progress iterator that yields items from the given iterable while
+    displaying a progress indicator in place on a single line of stderr.  The output is
+    not written to stdout, so the output of the program remains clean (see doctests).
+
+    for item in progress(range(1_000), desc="Processing"):
+        process(item)
+
+    Args:
+        items: The iterable of items to process.
+        desc: A description to display alongside the progress. Defaults to "".
+        total: The total number of items, defaults to 0. If 0, it will be inferred from
+            the iterable if possible.
+
+    Yields:
+        Iterator[T]: The items from the iterable, one by one.
+
+    >>> tuple(progress(range(5)))
+    (0, 1, 2, 3, 4)
+    >>> tuple(progress(range(5), desc="Processing", total=3))
+    (0, 1, 2, 3, 4)
+    >>> tuple(progress(range(5), desc="Processing", total=10))
+    (0, 1, 2, 3, 4)
+    >>> tuple(progress(range(5), desc="Processing", total=-5))
+    (0, 1, 2, 3, 4)
+    >>> from string import printable
+    >>> tuple(progress(printable, desc="Printable"))  # doctest: +ELLIPSIS
+    ('0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f',...
+    """
+    total = max(total, 0)
+    if not total and hasattr(items, "__len__"):
+        total = len(items)  # type: ignore[invalid-argument-type]
+
+    for i, item in enumerate(items, 1):
+        suffix = f"{i:,}/{total:,}" if total else f"{i:,}"
+        print(f"\r\033[K{desc}: {suffix}", end="", file=stderr, flush=True)
+        yield item
+
+    print("\r\033[K", end="", file=stderr, flush=True)
+
+
+if __name__ == "__main__":
+    import time
+
+    print("start")
+    for _item in progress(range(1_000), desc="Processing"):
+        time.sleep(0.02)
+
+    print("stop")

other/token_bucket.py

@@ -0,0 +1,100 @@
+"""
+Implementation of the Token Bucket Algorithm
+Token `rate` is added to the bucket every `frequency` seconds.
+The bucket can hold tokens up to `capacity` (full).
+The bucket starts full.
+Each request consumes one token.
+If a token arrives when the bucket is full, the token is discarded.
+If a request arrives when the bucket is empty, it is discarded.
+If the bucket has tokens available, requests will pass.
+https://en.wikipedia.org/wiki/Token_bucket
+"""
+
+import threading
+import time
+
+
+class TokenBucketRateLimiter:
+    def __init__(self, rate: int, capacity: int, frequency: int) -> None:
+        """
+        Initialize a Token Bucket rate limiter.
+
+        :param rate: Number of tokens added to the bucket per refill
+        :param capacity: Maximum number of tokens the bucket can hold.
+        :param frequency: Frequency of refill in seconds
+        >>> bucket = TokenBucketRateLimiter(4, 4, 60)
+        >>> bucket.tokens
+        4
+        >>> bucket.capacity
+        4
+        >>> bucket.frequency
+        60
+        """
+        self.rate = rate  # Tokens added per refill
+        self.capacity = capacity  # Maximum capacity of the bucket
+        self.frequency = frequency  # Frequency tokens are refilled
+        self.tokens = capacity  # Current tokens in the bucket
+        self.last_checked = time.time()  # Time when tokens were last checked
+        self.lock = threading.Lock()  # To make the rate limiter thread-safe
+
+    def _add_tokens(self) -> None:
+        """
+        Refill tokens only when a full minute has passed.
+        >>> bucket = TokenBucketRateLimiter(1, 4, 60)
+        >>> bucket.tokens  # Initially has a rate of 4 tokens
+        4
+        >>> bucket._add_tokens()
+        >>> bucket.tokens  # Bucket already full
+        4
+        """
+        current_time = time.time()
+        elapsed_time = current_time - self.last_checked
+
+        if elapsed_time >= self.frequency:
+            minutes_passed = int(elapsed_time // self.frequency)
+
+            # Add tokens based on rate
+            added_tokens = minutes_passed * self.rate
+            self.tokens = min(self.capacity, self.tokens + added_tokens)
+
+            # Update the last checked time
+            self.last_checked += minutes_passed * self.frequency
+
+    def allow_request(self) -> bool:
+        """
+        Check if a request is allowed.
+        If there are enough tokens, it consumes one token.
+        :return: True if the request is allowed, False otherwise.
+        >>> bucket = TokenBucketRateLimiter(1, 2, 60)
+        >>> bucket.allow_request()  # Token is available, request passes
+        True
+        >>> bucket.allow_request()  # Token is available, request passes
+        True
+        >>> bucket.allow_request()  # No token left, request is dropped
+        False
+        """
+        with self.lock:
+            self._add_tokens()
+            if self.tokens >= 1:
+                self.tokens -= 1
+                return True
+            return False
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()
+
+    print("Allow 4 requests per minute, capacity of 4")
+    bucket = TokenBucketRateLimiter(4, 4, 60)
+    total_requests = 10
+    delay_in_seconds = 10
+    print("Simulate 1 request per 10 seconds...")
+    for i in range(total_requests):
+        result = "pass" if bucket.allow_request() else "dropped"
+        print(
+            f"Request {i + 1}/{total_requests} \
+            timeline: {i * delay_in_seconds} seconds = {result}"
+        )
+        time.sleep(delay_in_seconds)