from time import time class TokenBucket(object): """An implementation of the token bucket algorithm. >>> bucket = TokenBucket(80, 0.5) >>> print bucket.consume(10) True >>> print bucket.consume(90) False """ def __init__(self, tokens, fill_rate): """tokens is the total tokens in the bucket. fill_rate is the rate in tokens/second that the bucket will be refilled.""" self.capacity = float(tokens) self._tokens = float(tokens) self.fill_rate = float(fill_rate) self.timestamp = time() def consume(self, tokens): """Consume tokens from the bucket. Returns True if there were sufficient tokens otherwise False.""" if tokens <= self.tokens: self._tokens -= tokens else: return False return True def refill(self): self._tokens = self.capacity def get_tokens(self): now = time() if self._tokens < self.capacity: delta = self.fill_rate * (now - self.timestamp) self._tokens = min(self.capacity, self._tokens + delta) self.timestamp = now return self._tokens tokens = property(get_tokens)