prototype

util/bucket.py

@@ -0,0 +1,40 @@
+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)