Spaces:
Sleeping
Sleeping
File size: 2,061 Bytes
c7dcc92 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
"""
Usage tracking module for rate limiting API requests.
This module provides a simple in-memory usage tracker that limits
the number of requests per user per day.
"""
from datetime import datetime
from typing import Dict
class UsageTracker:
"""Track and limit user requests on a daily basis."""
def __init__(self, daily_limit: int = 100):
"""
Initialize the usage tracker.
Args:
daily_limit: Maximum number of requests per user per day
"""
self.daily_limit = daily_limit
self.usage: Dict[datetime.date, Dict[str, int]] = {}
def check_limit(self, user_id: str) -> bool:
"""
Check if user has exceeded their daily limit and increment counter.
Args:
user_id: Unique identifier for the user (typically IP address)
Returns:
True if request is allowed, False if limit exceeded
"""
today = datetime.now().date()
# Clean up old dates to prevent memory growth
if today not in self.usage:
self.usage = {today: {}}
# Get current usage count for this user
user_count = self.usage[today].get(user_id, 0)
# Check if limit exceeded
if user_count >= self.daily_limit:
return False
# Increment counter
self.usage[today][user_id] = user_count + 1
return True
def get_usage(self, user_id: str) -> int:
"""
Get current usage count for a user today.
Args:
user_id: Unique identifier for the user
Returns:
Number of requests made today
"""
today = datetime.now().date()
return self.usage.get(today, {}).get(user_id, 0)
def get_remaining(self, user_id: str) -> int:
"""
Get remaining requests for a user today.
Args:
user_id: Unique identifier for the user
Returns:
Number of requests remaining today
"""
return max(0, self.daily_limit - self.get_usage(user_id))
|