usage_tracker.py

"""
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))