rate_limiter.py

from time import time


class TokenBucket(object):
    """An implementation of the token bucket algorithm.
    Usage example:
        from time import sleep
        import random

        t_0 = time()
        bucket = TokenBucket(2, 1)

        def take(n):
            print(f"Time: {time() - t_0}, "
                  f"Consume({n}) = {bucket.consume(n)}, "
                  f"Tokens = {bucket.get_level()}")

        while True:
            take(1)
            sleep(random.random() * 2)

    Recogniton to http://code.activestate.com/recipes/511490-implementation-of-the-token-bucket-algorithm/?in=lang-python
    Not thread safe
    """
    
    def __init__(self, capacity, fill_rate):
        """capacity is the total tokens in the bucket. fill_rate is the
        rate in tokens/second that the bucket will be refilled."""
        self.capacity = float(capacity)
        self.level = float(capacity)
        self.fill_rate = float(fill_rate)
        self.timestamp = time()

    def consume(self, units):
        """Consume tokens from the bucket. Returns True if there were
        sufficient tokens otherwise False."""
        if units <= self.get_level():
            self.level -= units
            return True
        else:
            return False

    def get_level(self):
        """Update the bucket level"""
        if self.level < self.capacity:
            now = time()
            delta = self.fill_rate * (now - self.timestamp)
            self.level = min(self.capacity, self.level + delta)
            self.timestamp = now
        return self.level