40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
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) |