sheriff-limiter 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,141 @@
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[cod]
4
+ *$py.class
5
+
6
+ # C extensions
7
+ *.so
8
+
9
+ # Distribution / packaging
10
+ bin/
11
+ build/
12
+ develop-eggs/
13
+ dist/
14
+ downloads/
15
+ eggs/
16
+ .eggs/
17
+ lib/
18
+ lib64/
19
+ parts/
20
+ sdist/
21
+ var/
22
+ wheels/
23
+ share/python-wheels/
24
+ *.egg-info/
25
+ .installed.cfg
26
+ *.egg
27
+ MANIFEST
28
+
29
+ # PyInstaller
30
+ # Usually these files are written by a python script from a template
31
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
32
+ *.manifest
33
+ *.spec
34
+
35
+ # Installer logs
36
+ pip-log.txt
37
+ pip-delete-this-directory.txt
38
+
39
+ # Unit test / coverage reports
40
+ htmlcov/
41
+ .tox/
42
+ .nox/
43
+ .coverage
44
+ .coverage.*
45
+ .cache
46
+ nosetests.xml
47
+ coverage.xml
48
+ *.cover
49
+ *.py,cover
50
+ .htmlcov/
51
+ .pytest_cache/
52
+ .ruff_cache/
53
+ .mypy_cache/
54
+
55
+ # Translations
56
+ *.mo
57
+ *.pot
58
+
59
+ # Django stuff:
60
+ *.log
61
+ local_settings.py
62
+ db.sqlite3
63
+ db.sqlite3-journal
64
+
65
+ # Sphinx documentation
66
+ doc/_build/
67
+
68
+ # PyBuilder
69
+ .pybuilder/
70
+ target/
71
+
72
+ # Jupyter Notebook
73
+ .ipynb_checkpoints
74
+
75
+ # IPython
76
+ profile_default/
77
+ ipython_config.py
78
+
79
+ # pyenv
80
+ # For a library or app, you might want to share your .python-version.
81
+ # Comment/uncomment to share.
82
+ #.python-version
83
+
84
+ # pipenv
85
+ # According to pypa/pipenv#1402, Pipfile.lock prevents deterministic
86
+ # builds in applications but is recommended for libraries.
87
+ # So, for a library you should keep it.
88
+ #Pipfile.lock
89
+
90
+ # poetry
91
+ # Similarly, poetry.lock should be committed for applications, but is optional/not recommended for libraries.
92
+ #poetry.lock
93
+
94
+ # pdm
95
+ # Similar to Pipfile.lock, pdm.lock is usage-specific.
96
+ #pdm.lock
97
+
98
+ # PEP 582; project local packages directory (used by pdm)
99
+ __pypackages__/
100
+
101
+ # Celery stuff
102
+ celerybeat-schedule
103
+ celerybeat.pid
104
+
105
+ # SageMath parsed files
106
+ *.sage.py
107
+
108
+ # Environments
109
+ .env
110
+ .venv
111
+ env/
112
+ venv/
113
+ ENV/
114
+ env.bak/
115
+ venv.bak/
116
+
117
+ # Spyder project settings
118
+ .spyderproject
119
+ .spyproject
120
+
121
+ # Rope project settings
122
+ .ropeproject
123
+
124
+ # mkdocs documentation
125
+ /site
126
+
127
+ # mypy
128
+ .mypy_cache/
129
+ .nosetests
130
+ /nosetests.xml
131
+ /.pytest_cache
132
+ .coverage
133
+
134
+ # Cython debug symbols
135
+ cython_debug/
136
+
137
+ # IDE files
138
+ .idea/
139
+ .vscode/
140
+ *.swp
141
+ *.swo
@@ -0,0 +1,153 @@
1
+ Metadata-Version: 2.4
2
+ Name: sheriff-limiter
3
+ Version: 0.1.0
4
+ Summary: An elegant, thread-safe, in-memory rate limiter for Python
5
+ Author-email: Vahsi Bati <info@vahsibati.com.tr>
6
+ License: MIT
7
+ Classifier: Development Status :: 3 - Alpha
8
+ Classifier: Intended Audience :: Developers
9
+ Classifier: License :: OSI Approved :: MIT License
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Programming Language :: Python :: 3.8
12
+ Classifier: Programming Language :: Python :: 3.9
13
+ Classifier: Programming Language :: Python :: 3.10
14
+ Classifier: Programming Language :: Python :: 3.11
15
+ Classifier: Programming Language :: Python :: 3.12
16
+ Requires-Python: >=3.8
17
+ Provides-Extra: dev
18
+ Requires-Dist: mypy>=1.0; extra == 'dev'
19
+ Requires-Dist: pytest-cov>=4.0; extra == 'dev'
20
+ Requires-Dist: pytest>=7.0; extra == 'dev'
21
+ Requires-Dist: ruff>=0.1.0; extra == 'dev'
22
+ Description-Content-Type: text/markdown
23
+
24
+ # Sheriff 🤠
25
+
26
+ An elegant, thread-safe, in-memory rate limiter for Python.
27
+
28
+ `sheriff` implements the **Token Bucket** algorithm, ensuring complete thread-safety with fine-grained locking and zero-leak memory management. It is designed to be lightweight, dependency-free, and extremely easy to integrate into any application or web framework (like FastAPI).
29
+
30
+ ---
31
+
32
+ ## Features
33
+
34
+ - 🔒 **Thread-Safe**: Uses fine-grained concurrent locks to ensure rate-limiting consistency across multiple threads.
35
+ - 🪣 **Token Bucket Algorithm**: Standard token bucket rate limiting with lazy, high-precision token replenishment.
36
+ - 🧹 **Self-Cleaning (Lazy Cleanup)**: Prunes stale/fully-replenished buckets from memory automatically to prevent memory leaks.
37
+ - ⚡ **Zero Dependencies**: Pure Python, built using standard library tools.
38
+ - 🚀 **FastAPI / Web Ready**: Fits perfectly into FastAPI's dependency injection (`Depends`) system.
39
+
40
+ ---
41
+
42
+ ## Installation
43
+
44
+ Install using `pip`:
45
+
46
+ ```bash
47
+ pip install sheriff-limiter
48
+ ```
49
+
50
+ ---
51
+
52
+ ## Quick Start
53
+
54
+ ### Basic Usage
55
+
56
+ Use `is_allowed` for a simple boolean check:
57
+
58
+ ```python
59
+ from sheriff import RateLimiter
60
+
61
+ # Default: 10 requests capacity, replenishes 1 token per second
62
+ limiter = RateLimiter()
63
+
64
+ # Check if allowed
65
+ if limiter.is_allowed("user_ip_address"):
66
+ print("Request allowed!")
67
+ else:
68
+ print("Rate limit exceeded.")
69
+ ```
70
+
71
+ ### Configuration Options
72
+
73
+ Initialize the limiter with custom parameters:
74
+
75
+ ```python
76
+ from sheriff import RateLimiter
77
+
78
+ # Configured for max 100 requests per minute
79
+ limiter = RateLimiter(max_requests=100, period=60.0)
80
+
81
+ # Or set capacity and refill rate directly
82
+ # Capacity of 5 tokens, refilling 0.5 tokens/sec
83
+ limiter = RateLimiter(capacity=5.0, refill_rate=0.5)
84
+ ```
85
+
86
+ ---
87
+
88
+ ## Advanced Features
89
+
90
+ ### 1. Raising Exceptions on Exceeding Limits
91
+
92
+ You can use `.check()` which raises a `RateLimitExceeded` exception. The exception contains a `retry_after` parameter telling you how long to wait in seconds.
93
+
94
+ ```python
95
+ from sheriff import RateLimiter, RateLimitExceeded
96
+
97
+ limiter = RateLimiter(max_requests=5, period=10.0)
98
+
99
+ try:
100
+ # Consume 1 token
101
+ limiter.check("client_1")
102
+ except RateLimitExceeded as e:
103
+ print(f"Rate limit exceeded! Retry after {e.retry_after:.2f} seconds.")
104
+ ```
105
+
106
+ ### 2. Manual Resets
107
+
108
+ Clear specific keys or reset all rate limits entirely:
109
+
110
+ ```python
111
+ # Reset a single client
112
+ limiter.reset("client_1")
113
+
114
+ # Reset all clients and clear the memory cache
115
+ limiter.reset_all()
116
+ ```
117
+
118
+ ---
119
+
120
+ ## FastAPI Integration
121
+
122
+ `sheriff` is perfect for FastAPI dependencies. Here is how you can use it to rate-limit endpoints by IP address:
123
+
124
+ ```python
125
+ from fastapi import FastAPI, Depends, Request, HTTPException, status
126
+ from sheriff import RateLimiter, RateLimitExceeded
127
+
128
+ app = FastAPI()
129
+
130
+ # 100 requests per minute limit
131
+ limiter = RateLimiter(max_requests=100, period=60.0)
132
+
133
+ def rate_limit(request: Request):
134
+ client_ip = request.client.host if request.client else "unknown"
135
+ try:
136
+ limiter.check(client_ip)
137
+ except RateLimitExceeded as e:
138
+ raise HTTPException(
139
+ status_code=status.HTTP_429_TOO_MANY_REQUESTS,
140
+ detail="Too many requests. Please slow down.",
141
+ headers={"Retry-After": str(int(e.retry_after or 0))}
142
+ )
143
+
144
+ @app.get("/items", dependencies=[Depends(rate_limit)])
145
+ async def read_items():
146
+ return {"status": "ok"}
147
+ ```
148
+
149
+ ---
150
+
151
+ ## License
152
+
153
+ MIT License. See [LICENSE](LICENSE) for details.
@@ -0,0 +1,130 @@
1
+ # Sheriff 🤠
2
+
3
+ An elegant, thread-safe, in-memory rate limiter for Python.
4
+
5
+ `sheriff` implements the **Token Bucket** algorithm, ensuring complete thread-safety with fine-grained locking and zero-leak memory management. It is designed to be lightweight, dependency-free, and extremely easy to integrate into any application or web framework (like FastAPI).
6
+
7
+ ---
8
+
9
+ ## Features
10
+
11
+ - 🔒 **Thread-Safe**: Uses fine-grained concurrent locks to ensure rate-limiting consistency across multiple threads.
12
+ - 🪣 **Token Bucket Algorithm**: Standard token bucket rate limiting with lazy, high-precision token replenishment.
13
+ - 🧹 **Self-Cleaning (Lazy Cleanup)**: Prunes stale/fully-replenished buckets from memory automatically to prevent memory leaks.
14
+ - ⚡ **Zero Dependencies**: Pure Python, built using standard library tools.
15
+ - 🚀 **FastAPI / Web Ready**: Fits perfectly into FastAPI's dependency injection (`Depends`) system.
16
+
17
+ ---
18
+
19
+ ## Installation
20
+
21
+ Install using `pip`:
22
+
23
+ ```bash
24
+ pip install sheriff-limiter
25
+ ```
26
+
27
+ ---
28
+
29
+ ## Quick Start
30
+
31
+ ### Basic Usage
32
+
33
+ Use `is_allowed` for a simple boolean check:
34
+
35
+ ```python
36
+ from sheriff import RateLimiter
37
+
38
+ # Default: 10 requests capacity, replenishes 1 token per second
39
+ limiter = RateLimiter()
40
+
41
+ # Check if allowed
42
+ if limiter.is_allowed("user_ip_address"):
43
+ print("Request allowed!")
44
+ else:
45
+ print("Rate limit exceeded.")
46
+ ```
47
+
48
+ ### Configuration Options
49
+
50
+ Initialize the limiter with custom parameters:
51
+
52
+ ```python
53
+ from sheriff import RateLimiter
54
+
55
+ # Configured for max 100 requests per minute
56
+ limiter = RateLimiter(max_requests=100, period=60.0)
57
+
58
+ # Or set capacity and refill rate directly
59
+ # Capacity of 5 tokens, refilling 0.5 tokens/sec
60
+ limiter = RateLimiter(capacity=5.0, refill_rate=0.5)
61
+ ```
62
+
63
+ ---
64
+
65
+ ## Advanced Features
66
+
67
+ ### 1. Raising Exceptions on Exceeding Limits
68
+
69
+ You can use `.check()` which raises a `RateLimitExceeded` exception. The exception contains a `retry_after` parameter telling you how long to wait in seconds.
70
+
71
+ ```python
72
+ from sheriff import RateLimiter, RateLimitExceeded
73
+
74
+ limiter = RateLimiter(max_requests=5, period=10.0)
75
+
76
+ try:
77
+ # Consume 1 token
78
+ limiter.check("client_1")
79
+ except RateLimitExceeded as e:
80
+ print(f"Rate limit exceeded! Retry after {e.retry_after:.2f} seconds.")
81
+ ```
82
+
83
+ ### 2. Manual Resets
84
+
85
+ Clear specific keys or reset all rate limits entirely:
86
+
87
+ ```python
88
+ # Reset a single client
89
+ limiter.reset("client_1")
90
+
91
+ # Reset all clients and clear the memory cache
92
+ limiter.reset_all()
93
+ ```
94
+
95
+ ---
96
+
97
+ ## FastAPI Integration
98
+
99
+ `sheriff` is perfect for FastAPI dependencies. Here is how you can use it to rate-limit endpoints by IP address:
100
+
101
+ ```python
102
+ from fastapi import FastAPI, Depends, Request, HTTPException, status
103
+ from sheriff import RateLimiter, RateLimitExceeded
104
+
105
+ app = FastAPI()
106
+
107
+ # 100 requests per minute limit
108
+ limiter = RateLimiter(max_requests=100, period=60.0)
109
+
110
+ def rate_limit(request: Request):
111
+ client_ip = request.client.host if request.client else "unknown"
112
+ try:
113
+ limiter.check(client_ip)
114
+ except RateLimitExceeded as e:
115
+ raise HTTPException(
116
+ status_code=status.HTTP_429_TOO_MANY_REQUESTS,
117
+ detail="Too many requests. Please slow down.",
118
+ headers={"Retry-After": str(int(e.retry_after or 0))}
119
+ )
120
+
121
+ @app.get("/items", dependencies=[Depends(rate_limit)])
122
+ async def read_items():
123
+ return {"status": "ok"}
124
+ ```
125
+
126
+ ---
127
+
128
+ ## License
129
+
130
+ MIT License. See [LICENSE](LICENSE) for details.
@@ -0,0 +1,63 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "sheriff-limiter"
7
+ dynamic = ["version"]
8
+ description = "An elegant, thread-safe, in-memory rate limiter for Python"
9
+ readme = "README.md"
10
+ requires-python = ">=3.8"
11
+ license = {text = "MIT"}
12
+ authors = [
13
+ {name = "Vahsi Bati", email = "info@vahsibati.com.tr"}
14
+ ]
15
+ classifiers = [
16
+ "Development Status :: 3 - Alpha",
17
+ "Intended Audience :: Developers",
18
+ "License :: OSI Approved :: MIT License",
19
+ "Programming Language :: Python :: 3",
20
+ "Programming Language :: Python :: 3.8",
21
+ "Programming Language :: Python :: 3.9",
22
+ "Programming Language :: Python :: 3.10",
23
+ "Programming Language :: Python :: 3.11",
24
+ "Programming Language :: Python :: 3.12",
25
+ ]
26
+ dependencies = []
27
+
28
+ [project.optional-dependencies]
29
+ dev = [
30
+ "pytest>=7.0",
31
+ "pytest-cov>=4.0",
32
+ "ruff>=0.1.0",
33
+ "mypy>=1.0",
34
+ ]
35
+
36
+ [tool.hatch.version]
37
+ path = "src/sheriff/__init__.py"
38
+
39
+ [tool.hatch.build.targets.wheel]
40
+ packages = ["src/sheriff"]
41
+
42
+ [tool.ruff]
43
+ line-length = 88
44
+ target-version = "py38"
45
+
46
+ [tool.ruff.lint]
47
+ select = [
48
+ "E", # pycodestyle errors
49
+ "W", # pycodestyle warnings
50
+ "F", # pyflakes
51
+ "I", # isort
52
+ "B", # flake8-bugbear
53
+ "C4", # flake8-comprehensions
54
+ "UP", # pyupgrade
55
+ ]
56
+ ignore = []
57
+
58
+ [tool.pytest.ini_options]
59
+ minversion = "7.0"
60
+ addopts = "-ra -q --tb=short"
61
+ testpaths = [
62
+ "tests",
63
+ ]
@@ -0,0 +1,10 @@
1
+ __version__ = "0.1.0"
2
+
3
+ from sheriff.core import RateLimiter
4
+ from sheriff.exceptions import RateLimitExceeded, SheriffError
5
+
6
+ __all__ = [
7
+ "RateLimiter",
8
+ "SheriffError",
9
+ "RateLimitExceeded",
10
+ ]
@@ -0,0 +1,217 @@
1
+ import time
2
+ from threading import Lock
3
+ from typing import Dict, Optional, Tuple
4
+
5
+ from sheriff.exceptions import RateLimitExceeded
6
+
7
+
8
+ class TokenBucket:
9
+ """Represents a single Token Bucket for rate limiting a specific key."""
10
+
11
+ def __init__(self, capacity: float, refill_rate: float):
12
+ self.capacity = capacity
13
+ self.refill_rate = refill_rate
14
+ self.tokens = capacity
15
+ self.last_updated = time.monotonic()
16
+ self.lock = Lock()
17
+
18
+ def consume(self, tokens: float = 1.0) -> Tuple[bool, float]:
19
+ """Consume tokens from the bucket.
20
+
21
+ Returns:
22
+ Tuple[bool, float]: (allowed, retry_after)
23
+ where allowed is True if consumed, False otherwise.
24
+ retry_after is the number of seconds to wait before there's enough tokens.
25
+ """
26
+ with self.lock:
27
+ now = time.monotonic()
28
+ elapsed = now - self.last_updated
29
+ if elapsed > 0:
30
+ self.tokens = min(
31
+ self.capacity, self.tokens + elapsed * self.refill_rate
32
+ )
33
+ self.last_updated = now
34
+
35
+ if self.tokens >= tokens:
36
+ self.tokens -= tokens
37
+ return True, 0.0
38
+
39
+ needed = tokens - self.tokens
40
+ retry_after = needed / self.refill_rate
41
+ return False, retry_after
42
+
43
+ def get_tokens(self) -> float:
44
+ """Returns the current number of tokens in the bucket after replenishment."""
45
+ with self.lock:
46
+ now = time.monotonic()
47
+ elapsed = now - self.last_updated
48
+ if elapsed > 0:
49
+ self.tokens = min(
50
+ self.capacity, self.tokens + elapsed * self.refill_rate
51
+ )
52
+ self.last_updated = now
53
+ return self.tokens
54
+
55
+ def reset(self) -> None:
56
+ """Resets the bucket to its full capacity."""
57
+ with self.lock:
58
+ self.tokens = self.capacity
59
+ self.last_updated = time.monotonic()
60
+
61
+ def is_full(self, now: float) -> bool:
62
+ """Check if the bucket is fully replenished.
63
+ Must be called under RateLimiter container lock or self.lock.
64
+ """
65
+ with self.lock:
66
+ elapsed = now - self.last_updated
67
+ current_tokens = min(
68
+ self.capacity, self.tokens + elapsed * self.refill_rate
69
+ )
70
+ return current_tokens >= self.capacity
71
+
72
+
73
+ class RateLimiter:
74
+ """Core class representing the Sheriff thread-safe, in-memory rate limiter."""
75
+
76
+ def __init__(
77
+ self,
78
+ capacity: float = 10.0,
79
+ refill_rate: float = 1.0,
80
+ max_requests: Optional[int] = None,
81
+ period: Optional[float] = None,
82
+ cleanup_interval: float = 60.0,
83
+ ):
84
+ """Initializes the RateLimiter.
85
+
86
+ Args:
87
+ capacity: Maximum number of tokens a bucket can hold. Defaults to 10.0.
88
+ refill_rate: Number of tokens added to the bucket per second.
89
+ Defaults to 1.0.
90
+ max_requests: Optional parameter to initialize capacity using requests.
91
+ period: Optional parameter to specify the period in seconds for
92
+ max_requests.
93
+ cleanup_interval: Time in seconds between periodic cleanup sweeps of
94
+ fully replenished buckets. Defaults to 60.0.
95
+ """
96
+ if max_requests is not None:
97
+ capacity = float(max_requests)
98
+ if period is not None:
99
+ refill_rate = capacity / period
100
+ else:
101
+ refill_rate = capacity
102
+
103
+ if capacity <= 0:
104
+ raise ValueError("Capacity must be greater than zero.")
105
+ if refill_rate <= 0:
106
+ raise ValueError("Refill rate must be greater than zero.")
107
+ if cleanup_interval <= 0:
108
+ raise ValueError("Cleanup interval must be greater than zero.")
109
+
110
+ self.capacity = capacity
111
+ self.refill_rate = refill_rate
112
+ self.cleanup_interval = cleanup_interval
113
+
114
+ self.buckets: Dict[str, TokenBucket] = {}
115
+ self.lock = Lock()
116
+ self.last_cleanup = time.monotonic()
117
+
118
+ def _get_bucket(self, key: str) -> TokenBucket:
119
+ """Thread-safe retrieval or creation of a TokenBucket for a given key.
120
+ Also triggers lazy cleanup if the cleanup interval has elapsed.
121
+ """
122
+ with self.lock:
123
+ self._maybe_cleanup()
124
+ if key not in self.buckets:
125
+ self.buckets[key] = TokenBucket(self.capacity, self.refill_rate)
126
+ return self.buckets[key]
127
+
128
+ def _maybe_cleanup(self) -> None:
129
+ """Prunes fully replenished buckets from memory.
130
+ Must be called with self.lock held.
131
+ """
132
+ now = time.monotonic()
133
+ if now - self.last_cleanup >= self.cleanup_interval:
134
+ keys_to_delete = []
135
+ for key, bucket in self.buckets.items():
136
+ if bucket.is_full(now):
137
+ keys_to_delete.append(key)
138
+ for key in keys_to_delete:
139
+ del self.buckets[key]
140
+ self.last_cleanup = now
141
+
142
+ def is_allowed(self, key: str, tokens: float = 1.0) -> bool:
143
+ """Check if the request is allowed under the rate limit.
144
+
145
+ Args:
146
+ key: Unique identifier for the client or bucket.
147
+ tokens: Number of tokens to consume. Defaults to 1.0.
148
+
149
+ Returns:
150
+ bool: True if the request is allowed, False otherwise.
151
+ """
152
+ bucket = self._get_bucket(key)
153
+ allowed, _ = bucket.consume(tokens)
154
+ return allowed
155
+
156
+ def check(self, key: str, tokens: float = 1.0) -> None:
157
+ """Check if the request is allowed under the rate limit.
158
+ Raises RateLimitExceeded if not.
159
+
160
+ Args:
161
+ key: Unique identifier for the client or bucket.
162
+ tokens: Number of tokens to consume. Defaults to 1.0.
163
+
164
+ Raises:
165
+ RateLimitExceeded: If the key is rate-limited.
166
+ """
167
+ bucket = self._get_bucket(key)
168
+ allowed, retry_after = bucket.consume(tokens)
169
+ if not allowed:
170
+ raise RateLimitExceeded(
171
+ message=f"Rate limit exceeded for key: {key}",
172
+ retry_after=retry_after,
173
+ )
174
+
175
+ def consume(self, key: str, tokens: float = 1.0) -> Tuple[bool, float]:
176
+ """Consume tokens from the bucket for the given key.
177
+
178
+ Args:
179
+ key: Unique identifier for the client or bucket.
180
+ tokens: Number of tokens to consume. Defaults to 1.0.
181
+
182
+ Returns:
183
+ Tuple[bool, float]: (allowed, retry_after)
184
+ where allowed is True if consumed, False otherwise.
185
+ retry_after is the number of seconds to wait before there are
186
+ enough tokens.
187
+ """
188
+ bucket = self._get_bucket(key)
189
+ return bucket.consume(tokens)
190
+
191
+ def get_tokens(self, key: str) -> float:
192
+ """Returns the current number of tokens available in the bucket for the key.
193
+
194
+ Args:
195
+ key: Unique identifier for the client or bucket.
196
+
197
+ Returns:
198
+ float: Current number of tokens.
199
+ """
200
+ bucket = self._get_bucket(key)
201
+ return bucket.get_tokens()
202
+
203
+ def reset(self, key: str) -> None:
204
+ """Reset the rate limit bucket for the given key.
205
+
206
+ Args:
207
+ key: Unique identifier for the client or bucket.
208
+ """
209
+ with self.lock:
210
+ bucket = self.buckets.get(key)
211
+ if bucket is not None:
212
+ bucket.reset()
213
+
214
+ def reset_all(self) -> None:
215
+ """Reset all rate limit buckets, clearing the internal cache."""
216
+ with self.lock:
217
+ self.buckets.clear()
@@ -0,0 +1,19 @@
1
+ from typing import Optional
2
+
3
+
4
+ class SheriffError(Exception):
5
+ """Base exception for all Sheriff rate limiter errors."""
6
+
7
+ pass
8
+
9
+
10
+ class RateLimitExceeded(SheriffError):
11
+ """Exception raised when a rate limit is exceeded."""
12
+
13
+ def __init__(
14
+ self,
15
+ message: str = "Rate limit exceeded",
16
+ retry_after: Optional[float] = None,
17
+ ):
18
+ super().__init__(message)
19
+ self.retry_after = retry_after
@@ -0,0 +1 @@
1
+ # Test suite for sheriff rate limiter
@@ -0,0 +1,9 @@
1
+ import pytest
2
+
3
+ from sheriff.core import RateLimiter
4
+
5
+
6
+ @pytest.fixture
7
+ def limiter():
8
+ """Returns a basic RateLimiter instance."""
9
+ return RateLimiter()
@@ -0,0 +1,193 @@
1
+ import threading
2
+ import time
3
+ from unittest.mock import patch
4
+
5
+ import pytest
6
+
7
+ from sheriff.core import RateLimiter
8
+ from sheriff.exceptions import RateLimitExceeded
9
+
10
+
11
+ def test_limiter_creation_defaults(limiter):
12
+ assert isinstance(limiter, RateLimiter)
13
+ assert limiter.capacity == 10.0
14
+ assert limiter.refill_rate == 1.0
15
+ assert limiter.cleanup_interval == 60.0
16
+
17
+
18
+ def test_limiter_creation_custom():
19
+ limiter = RateLimiter(capacity=20.0, refill_rate=2.5, cleanup_interval=30.0)
20
+ assert limiter.capacity == 20.0
21
+ assert limiter.refill_rate == 2.5
22
+ assert limiter.cleanup_interval == 30.0
23
+
24
+
25
+ def test_limiter_creation_max_requests():
26
+ limiter = RateLimiter(max_requests=100, period=60.0)
27
+ assert limiter.capacity == 100.0
28
+ assert limiter.refill_rate == 100.0 / 60.0
29
+
30
+
31
+ def test_limiter_creation_max_requests_no_period():
32
+ limiter = RateLimiter(max_requests=50)
33
+ assert limiter.capacity == 50.0
34
+ assert limiter.refill_rate == 50.0
35
+
36
+
37
+ def test_invalid_parameters():
38
+ with pytest.raises(ValueError, match="Capacity must be greater than zero"):
39
+ RateLimiter(capacity=0)
40
+ with pytest.raises(ValueError, match="Capacity must be greater than zero"):
41
+ RateLimiter(capacity=-10)
42
+ with pytest.raises(ValueError, match="Refill rate must be greater than zero"):
43
+ RateLimiter(refill_rate=0)
44
+ with pytest.raises(ValueError, match="Refill rate must be greater than zero"):
45
+ RateLimiter(refill_rate=-1)
46
+ with pytest.raises(ValueError, match="Cleanup interval must be greater than zero"):
47
+ RateLimiter(cleanup_interval=0)
48
+ with pytest.raises(ValueError, match="Cleanup interval must be greater than zero"):
49
+ RateLimiter(cleanup_interval=-5)
50
+
51
+
52
+ def test_is_allowed_basic():
53
+ limiter = RateLimiter(capacity=3.0, refill_rate=1.0)
54
+
55
+ assert limiter.is_allowed("user-1", tokens=1.0) is True
56
+ assert limiter.is_allowed("user-1", tokens=2.0) is True
57
+ # Now empty
58
+ assert limiter.is_allowed("user-1", tokens=1.0) is False
59
+
60
+
61
+ def test_token_replenishment():
62
+ limiter = RateLimiter(capacity=5.0, refill_rate=2.0)
63
+
64
+ start_time = 100.0
65
+ with patch("time.monotonic", return_value=start_time):
66
+ assert limiter.is_allowed("user-2", tokens=5.0) is True
67
+ assert limiter.is_allowed("user-2", tokens=1.0) is False
68
+
69
+ # After 1.5 seconds, we should replenish 1.5 * 2 = 3.0 tokens
70
+ with patch("time.monotonic", return_value=start_time + 1.5):
71
+ assert limiter.is_allowed("user-2", tokens=3.0) is True
72
+ assert limiter.is_allowed("user-2", tokens=1.0) is False
73
+
74
+ # After another 2.5 seconds, we should replenish up to capacity (max 5)
75
+ with patch("time.monotonic", return_value=start_time + 4.0):
76
+ assert limiter.is_allowed("user-2", tokens=5.0) is True
77
+ assert limiter.is_allowed("user-2", tokens=1.0) is False
78
+
79
+
80
+ def test_check_raises_exception():
81
+ limiter = RateLimiter(capacity=2.0, refill_rate=1.0)
82
+
83
+ start_time = 100.0
84
+ with patch("time.monotonic", return_value=start_time):
85
+ limiter.check("user-3", tokens=2.0)
86
+
87
+ with pytest.raises(RateLimitExceeded) as exc_info:
88
+ limiter.check("user-3", tokens=1.0)
89
+
90
+ assert exc_info.value.retry_after == 1.0
91
+
92
+ # If we need 2 tokens, it will require 2 seconds
93
+ with pytest.raises(RateLimitExceeded) as exc_info2:
94
+ limiter.check("user-3", tokens=2.0)
95
+
96
+ assert exc_info2.value.retry_after == 2.0
97
+
98
+
99
+ def test_consume_returns_tuple():
100
+ limiter = RateLimiter(capacity=2.0, refill_rate=0.5)
101
+
102
+ start_time = 100.0
103
+ with patch("time.monotonic", return_value=start_time):
104
+ allowed, retry_after = limiter.consume("user-4", tokens=1.5)
105
+ assert allowed is True
106
+ assert retry_after == 0.0
107
+
108
+ allowed, retry_after = limiter.consume("user-4", tokens=1.0)
109
+ assert allowed is False
110
+ # Needs 1.0 - 0.5 = 0.5 tokens. At refill_rate 0.5, needs 1.0 second.
111
+ assert retry_after == 1.0
112
+
113
+
114
+ def test_get_tokens():
115
+ limiter = RateLimiter(capacity=10.0, refill_rate=2.0)
116
+
117
+ start_time = 100.0
118
+ with patch("time.monotonic", return_value=start_time):
119
+ assert limiter.get_tokens("user-5") == 10.0
120
+ assert limiter.is_allowed("user-5", tokens=4.0) is True
121
+ assert limiter.get_tokens("user-5") == 6.0
122
+
123
+ with patch("time.monotonic", return_value=start_time + 1.5):
124
+ # 6.0 + 1.5 * 2 = 9.0
125
+ assert limiter.get_tokens("user-5") == 9.0
126
+
127
+
128
+ def test_reset_and_reset_all():
129
+ limiter = RateLimiter(capacity=5.0, refill_rate=1.0)
130
+
131
+ limiter.is_allowed("user-a", tokens=5.0)
132
+ limiter.is_allowed("user-b", tokens=5.0)
133
+
134
+ assert limiter.is_allowed("user-a", tokens=1.0) is False
135
+ assert limiter.is_allowed("user-b", tokens=1.0) is False
136
+
137
+ limiter.reset("user-a")
138
+ assert limiter.is_allowed("user-a", tokens=5.0) is True
139
+ assert limiter.is_allowed("user-b", tokens=1.0) is False
140
+
141
+ limiter.is_allowed("user-a", tokens=5.0)
142
+ limiter.reset_all()
143
+ assert limiter.is_allowed("user-a", tokens=5.0) is True
144
+ assert limiter.is_allowed("user-b", tokens=5.0) is True
145
+
146
+
147
+ def test_lazy_cleanup():
148
+ limiter = RateLimiter(capacity=5.0, refill_rate=1.0, cleanup_interval=0.1)
149
+
150
+ # Access a key to create a bucket
151
+ assert limiter.is_allowed("key1", tokens=1.0) is True
152
+ assert "key1" in limiter.buckets
153
+
154
+ # Wait, but not long enough to fully replenish
155
+ time.sleep(0.15)
156
+
157
+ # Access key2 to trigger cleanup check
158
+ assert limiter.is_allowed("key2", tokens=1.0) is True
159
+ # key1 is not fully replenished, so it shouldn't be deleted
160
+ assert "key1" in limiter.buckets
161
+
162
+ # Wait long enough to fully replenish key1 (needs 1.0 second to recover 1.0 token)
163
+ time.sleep(1.0)
164
+
165
+ # Access key2 to trigger cleanup
166
+ assert limiter.is_allowed("key2", tokens=1.0) is True
167
+ # key1 is now fully replenished, so it should be deleted
168
+ assert "key1" not in limiter.buckets
169
+
170
+
171
+ def test_concurrent_consumption():
172
+ # Thread safety check: Ensure no double-consumption
173
+ limiter = RateLimiter(capacity=50.0, refill_rate=0.0001)
174
+
175
+ successes = [0]
176
+ lock = threading.Lock()
177
+
178
+ def worker():
179
+ for _ in range(10):
180
+ if limiter.is_allowed("concurrent-key", tokens=1.0):
181
+ with lock:
182
+ successes[0] += 1
183
+
184
+ threads = [threading.Thread(target=worker) for _ in range(10)]
185
+ for t in threads:
186
+ t.start()
187
+ for t in threads:
188
+ t.join()
189
+
190
+ # With 10 threads doing 10 attempts each, total 100 attempts,
191
+ # but capacity is 50 and refill rate is negligible.
192
+ # Therefore, exactly 50 should succeed.
193
+ assert successes[0] == 50