Rate Limiting Utilities

from limits import parse, strategies, storage, RateLimitExceeded
import time
import os

# 1. Initialize a storage backend
# For simplicity, using in-memory storage. For production, use Redis/Memcached.
# A Redis example:
# REDIS_URL = os.environ.get('REDIS_URL', 'redis://localhost:6379')
# store = storage.RedisStorage(REDIS_URL)

store = storage.MemoryStorage()

# 2. Define a rate limit (e.g., 5 requests per minute)
# parse() converts a string like '5/minute' into a RateLimitItem
rate_limit_string = '5/minute'
rate_limit = parse(rate_limit_string)

# 3. Initialize a rate limiter strategy
# FixedWindowRateLimiter is a common strategy
limiter = strategies.FixedWindowRateLimiter(store)

def do_limited_action(user_id):
    try:
        # 4. Test the limit for a specific identifier (e.g., user_id)
        if limiter.test(rate_limit, user_id):
            # 5. Consume the limit (i.e., record a hit)
            limiter.hit(rate_limit, user_id)
            print(f"[{user_id}] Action performed at {time.strftime('%H:%M:%S')}. Remaining: {limiter.get_remaining(rate_limit, user_id)}")
        else:
            raise RateLimitExceeded(rate_limit)
    except RateLimitExceeded as e:
        # Query available capacity and reset time
        remaining = limiter.get_remaining(e.limit, user_id)
        reset_at = limiter.get_reset_time(e.limit, user_id)
        print(f"[{user_id}] Rate limit exceeded for {e.limit}. Try again in {reset_at - time.time():.1f} seconds. Remaining: {remaining}")

# Simulate some requests
user = "test_user_1"
print(f"--- Simulating requests for {user} ({rate_limit_string}) ---")
for i in range(7):
    do_limited_action(user)
    time.sleep(1) # Simulate some delay

print("\n--- Waiting for reset time ---")
time.sleep(60)

print("\n--- Simulating requests after reset ---")
for i in range(3):
    do_limited_action(user)
    time.sleep(1)