Merge pull request #93 from aaronfultonnz/main

Added feature to enable custom buckets

README.md

@@ -189,6 +189,18 @@ session.get('https://api.some_site.com/v1/')
 session.get('https://api.some_site.com/v1/users/1234')
 ```
 
+### Custom Tracking
+For advanced use cases, you can define your own custom tracking behavior with the `bucket` option.
+For example, an API that enforces rate limits based on a tenant ID, this feature can be used to track
+rate limits per tenant. If `bucket` is specified, host tracking is disabled.
+
+Note: It is advisable to use SQLite or Redis backends when using custom tracking because using the default backend
+each session will track rate limits independently, even if both sessions call the same URL.
+```python
+sessionA = LimiterSession(per_second=5, bucket='tenant1')
+sessionB = LimiterSession(per_second=5, bucket='tenant2')
+```
+
 ### Rate Limit Error Handling
 Sometimes, server-side rate limiting may not behave exactly as documented (or may not be documented
 at all). Or you might encounter other scenarios where your client-side limit gets out of sync with

requests_ratelimiter/requests_ratelimiter.py

@@ -39,6 +39,7 @@ def __init__(
         max_delay: Union[int, float, None] = None,
         per_host: bool = True,
         limit_statuses: Iterable[int] = (429,),
+        bucket_name: Optional[str] = None,
         **kwargs,
     ):
         # Translate request rate values into RequestRate objects
@@ -72,6 +73,7 @@ def __init__(
         self.limit_statuses = limit_statuses
         self.max_delay = max_delay
         self.per_host = per_host
+        self.bucket_name = bucket_name
         self._default_bucket = str(uuid4())
 
         # If the superclass is an adapter or custom Session, pass along any valid keyword arguments
@@ -97,7 +99,12 @@ def send(self, request: PreparedRequest, **kwargs) -> Response:
 
     def _bucket_name(self, request):
         """Get a bucket name for the given request"""
-        return urlparse(request.url).netloc if self.per_host else self._default_bucket
+        if self.bucket_name:
+            return self.bucket_name
+        elif self.per_host:
+            return urlparse(request.url).netloc
+        else:
+            return self._default_bucket
 
     def _fill_bucket(self, request: PreparedRequest):
         """Partially fill the bucket for the given request, requiring an extra delay until the next

test/test_requests_ratelimiter.py

@@ -95,6 +95,19 @@ def test_custom_session(mock_sleep):
     session.get(MOCKED_URL)
     assert mock_sleep.called is True
 
+@patch_sleep
+def test_custom_bucket(mock_sleep):
+    """With custom buckets, each session can be called independently without triggering rate limiting"""
+    session_a = get_mock_session(per_second=5, bucket_name="a")
+    session_b = get_mock_session(per_second=5, bucket_name="b")
+
+    for _ in range(5):
+        session_a.get(MOCKED_URL)
+        session_b.get(MOCKED_URL)
+    assert mock_sleep.called is False    
+
+    session_a.get(MOCKED_URL)    
+    assert mock_sleep.called is True
 
 @patch_sleep
 def test_429(mock_sleep):