Switch to Google-style docs and implement hash in the token bucket (#50)

Author: kgaughan

src/uwhoisd/rl.py

@@ -1,13 +1,16 @@
-"""
-Rate limiting.
-"""
+"""Rate limiting."""
 
+import functools
 import time
 
 
+@functools.total_ordering
 class TokenBucket:
-    """
-    A token bucket.
+    """A token bucket.
+
+    Args:
+        rate: token refill rate (per second) of bucket
+        limit: maximum number of tokens the bucket can contain
     """
 
     __slots__ = [
@@ -19,25 +22,21 @@ class TokenBucket:
 
     clock = staticmethod(time.time)
 
-    def __init__(self, rate, limit):
-        """
-        Create a new token bucket.
-
-        :param rate int: Token refill rate (per second) of bucket.
-        :param limit int: Maximum number of tokens the bucket can contain.
-        """
+    def __init__(self, rate: int, limit: int) -> None:
         super().__init__()
         self.ts = self.clock()
         self.rate = rate
         self.limit = limit
         self._available = limit
 
-    def consume(self, tokens):
-        """
-        Attempt to remove the given number of tokens from this bucket.
+    def consume(self, tokens: int) -> bool:
+        """Attempt to remove the given number of tokens from this bucket.
+
+        Args:
+            tokens: number of tokens to consume
 
-        If there are enough tokens in the bucket to fulfil the request, we
-        return `True`, otherwise `False`.
+        Returns:
+            `True` if the requested number could be consumed, otherwise `False`.
         """
         if 0 <= tokens <= self.tokens:
             self._available -= tokens
@@ -46,8 +45,10 @@ def consume(self, tokens):
 
     @property
     def tokens(self):
-        """
-        Gives the number of tokens available in this bucket.
+        """Gives the number of tokens available in this bucket.
+
+        Returns:
+            The number of available tokens.
         """
         ts = self.clock()
         if self._available < self.limit:
@@ -61,6 +62,9 @@ def __eq__(self, other):
     def __lt__(self, other):
         return self.ts < other.ts
 
+    def __hash__(self):
+        return hash((self._available, self.limit, self.rate, self.ts))
+
     def __getstate__(self):
         return dict(zip(self.__slots__, [getattr(self, attr) for attr in self.__slots__]))
 

tests/utils.py

@@ -1,6 +1,4 @@
-"""
-Utility functions for testing.
-"""
+"""Utility functions for testing."""
 
 from os import path
 
@@ -11,9 +9,7 @@
 
 
 class Clock:
-    """
-    A fake clock.
-    """
+    """A fake clock."""
 
     def __init__(self, initial: int = 0):
         super().__init__()