1"""Token bucket implementation for rate limiting."""
2from __future__ import absolute_import, unicode_literals
3
4from collections import deque
5
6from kombu.five import monotonic
7
8__all__ = ('TokenBucket',)
9
10
11class TokenBucket(object):
12    """Token Bucket Algorithm.
13
14    See Also:
15        https://en.wikipedia.org/wiki/Token_Bucket
16
17        Most of this code was stolen from an entry in the ASPN Python Cookbook:
18        https://code.activestate.com/recipes/511490/
19
20    Warning:
21        Thread Safety: This implementation is not thread safe.
22        Access to a `TokenBucket` instance should occur within the critical
23        section of any multithreaded code.
24    """
25
26    #: The rate in tokens/second that the bucket will be refilled.
27    fill_rate = None
28
29    #: Maximum number of tokens in the bucket.
30    capacity = 1
31
32    #: Timestamp of the last time a token was taken out of the bucket.
33    timestamp = None
34
35    def __init__(self, fill_rate, capacity=1):
36        self.capacity = float(capacity)
37        self._tokens = capacity
38        self.fill_rate = float(fill_rate)
39        self.timestamp = monotonic()
40        self.contents = deque()
41
42    def add(self, item):
43        self.contents.append(item)
44
45    def pop(self):
46        return self.contents.popleft()
47
48    def clear_pending(self):
49        self.contents.clear()
50
51    def can_consume(self, tokens=1):
52        """Check if one or more tokens can be consumed.
53
54        Returns:
55            bool: true if the number of tokens can be consumed
56                from the bucket.  If they can be consumed, a call will also
57                consume the requested number of tokens from the bucket.
58                Calls will only consume `tokens` (the number requested)
59                or zero tokens -- it will never consume a partial number
60                of tokens.
61        """
62        if tokens <= self._get_tokens():
63            self._tokens -= tokens
64            return True
65        return False
66
67    def expected_time(self, tokens=1):
68        """Return estimated time of token availability.
69
70        Returns:
71            float: the time in seconds.
72        """
73        _tokens = self._get_tokens()
74        tokens = max(tokens, _tokens)
75        return (tokens - _tokens) / self.fill_rate
76
77    def _get_tokens(self):
78        if self._tokens < self.capacity:
79            now = monotonic()
80            delta = self.fill_rate * (now - self.timestamp)
81            self._tokens = min(self.capacity, self._tokens + delta)
82            self.timestamp = now
83        return self._tokens
84