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 get_tokens(self):
if self._tokens < self.capacity:
now = time()
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)
if __name__ == '__main__':
from time import sleep
bucket = TokenBucket(80, 1)
print "tokens =", bucket.tokens
print "consume(10) =", bucket.consume(10)
print "consume(10) =", bucket.consume(10)
sleep(1)
print "tokens =", bucket.tokens
sleep(1)
print "tokens =", bucket.tokens
print "consume(90) =", bucket.consume(90)
print "tokens =", bucket.tokens
