diff --git a/other/cheap_progress.py b/other/cheap_progress.py new file mode 100644 index 000000000..62da9b59d --- /dev/null +++ b/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") diff --git a/other/token_bucket.py b/other/token_bucket.py new file mode 100644 index 000000000..3bfa6ec79 --- /dev/null +++ b/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)