Skip to content

Commit da4622c

Browse files
committed
feat: add RateLimiter utility class
1 parent 76392b5 commit da4622c

File tree

3 files changed

+102
-0
lines changed

3 files changed

+102
-0
lines changed

src/gradient/_utils/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
get_required_header as get_required_header,
3030
maybe_coerce_boolean as maybe_coerce_boolean,
3131
maybe_coerce_integer as maybe_coerce_integer,
32+
RateLimiter as RateLimiter,
3233
)
3334
from ._compat import (
3435
get_args as get_args,

src/gradient/_utils/_utils.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -419,3 +419,48 @@ def json_safe(data: object) -> object:
419419
return data.isoformat()
420420

421421
return data
422+
423+
424+
# Rate Limiting Classes
425+
class RateLimiter:
426+
"""Simple token bucket rate limiter."""
427+
428+
def __init__(self, requests_per_minute: int = 60) -> None:
429+
"""Initialize rate limiter.
430+
431+
Args:
432+
requests_per_minute: Maximum requests allowed per minute
433+
"""
434+
self.requests_per_minute: int = requests_per_minute
435+
self.tokens: float = float(requests_per_minute)
436+
self.last_refill: float = self._now()
437+
self.refill_rate: float = requests_per_minute / 60.0 # tokens per second
438+
439+
def _now(self) -> float:
440+
"""Get current time in seconds."""
441+
import time
442+
return time.time()
443+
444+
def _refill(self) -> None:
445+
"""Refill tokens based on elapsed time."""
446+
now = self._now()
447+
elapsed = now - self.last_refill
448+
self.tokens = min(self.requests_per_minute, self.tokens + elapsed * self.refill_rate)
449+
self.last_refill = now
450+
451+
def acquire(self, tokens: int = 1) -> bool:
452+
"""Try to acquire tokens. Returns True if successful."""
453+
self._refill()
454+
if self.tokens >= tokens:
455+
self.tokens -= tokens
456+
return True
457+
return False
458+
459+
def wait_time(self, tokens: int = 1) -> float:
460+
"""Get seconds to wait for tokens to be available."""
461+
self._refill()
462+
if self.tokens >= tokens:
463+
return 0.0
464+
465+
needed = tokens - self.tokens
466+
return needed / self.refill_rate

tests/test_rate_limiter.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
"""Tests for rate limiting functionality."""
2+
3+
import time
4+
import pytest
5+
from gradient._utils import RateLimiter
6+
7+
8+
class TestRateLimiter:
9+
"""Test rate limiting functionality."""
10+
11+
def test_rate_limiter_basic(self):
12+
"""Test basic rate limiter operations."""
13+
limiter = RateLimiter(requests_per_minute=10)
14+
15+
# Should allow initial requests
16+
assert limiter.acquire() is True
17+
assert limiter.acquire() is True
18+
19+
# Should deny when tokens exhausted
20+
limiter.tokens = 0 # Force exhaustion
21+
assert limiter.acquire() is False
22+
23+
def test_rate_limiter_wait_time(self):
24+
"""Test wait time calculation."""
25+
limiter = RateLimiter(requests_per_minute=60) # 1 request per second
26+
27+
# Exhaust tokens
28+
limiter.tokens = 0
29+
30+
# Should calculate correct wait time
31+
wait_time = limiter.wait_time()
32+
assert wait_time > 0
33+
assert wait_time <= 1.0 # Should not exceed 1 second
34+
35+
def test_rate_limiter_refill(self):
36+
"""Test token refill over time."""
37+
limiter = RateLimiter(requests_per_minute=60) # 1 token per second
38+
39+
# Exhaust tokens
40+
limiter.tokens = 0
41+
start_time = limiter._now()
42+
43+
# Wait for refill
44+
time.sleep(0.1)
45+
46+
# Should have refilled some tokens
47+
limiter._refill()
48+
assert limiter.tokens > 0
49+
50+
def test_rate_limiter_custom_rate(self):
51+
"""Test custom rate limits."""
52+
limiter = RateLimiter(requests_per_minute=120) # 2 requests per second
53+
54+
# Should have double the tokens of default
55+
assert limiter.requests_per_minute == 120
56+
assert limiter.refill_rate == 2.0

0 commit comments

Comments
 (0)