fix: auto-fix code issues (cron)
- 修复重复导入/字段 - 修复异常处理 - 修复PEP8格式问题 - 添加类型注解 - 修复缺失的urllib.parse导入
This commit is contained in:
@@ -12,6 +12,7 @@ from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from functools import wraps
|
||||
|
||||
|
||||
@dataclass
|
||||
class RateLimitConfig:
|
||||
"""限流配置"""
|
||||
@@ -20,6 +21,7 @@ class RateLimitConfig:
|
||||
burst_size: int = 10 # 突发请求数
|
||||
window_size: int = 60 # 窗口大小(秒)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RateLimitInfo:
|
||||
"""限流信息"""
|
||||
@@ -29,6 +31,7 @@ class RateLimitInfo:
|
||||
reset_time: int # 重置时间戳
|
||||
retry_after: int # 需要等待的秒数
|
||||
|
||||
|
||||
class SlidingWindowCounter:
|
||||
"""滑动窗口计数器"""
|
||||
|
||||
@@ -60,6 +63,7 @@ class SlidingWindowCounter:
|
||||
for k in old_keys:
|
||||
self.requests.pop(k, None)
|
||||
|
||||
|
||||
class RateLimiter:
|
||||
"""API 限流器"""
|
||||
|
||||
@@ -106,13 +110,18 @@ class RateLimiter:
|
||||
# 检查是否超过限制
|
||||
if current_count >= stored_config.requests_per_minute:
|
||||
return RateLimitInfo(
|
||||
allowed=False, remaining=0, reset_time=reset_time, retry_after=stored_config.window_size
|
||||
allowed=False,
|
||||
remaining=0,
|
||||
reset_time=reset_time,
|
||||
retry_after=stored_config.window_size,
|
||||
)
|
||||
|
||||
# 允许请求,增加计数
|
||||
await counter.add_request()
|
||||
|
||||
return RateLimitInfo(allowed=True, remaining=remaining - 1, reset_time=reset_time, retry_after=0)
|
||||
return RateLimitInfo(
|
||||
allowed=True, remaining=remaining - 1, reset_time=reset_time, retry_after=0
|
||||
)
|
||||
|
||||
async def get_limit_info(self, key: str) -> RateLimitInfo:
|
||||
"""获取限流信息(不增加计数)"""
|
||||
@@ -136,7 +145,9 @@ class RateLimiter:
|
||||
allowed=current_count < config.requests_per_minute,
|
||||
remaining=remaining,
|
||||
reset_time=reset_time,
|
||||
retry_after=max(0, config.window_size) if current_count >= config.requests_per_minute else 0,
|
||||
retry_after=max(0, config.window_size)
|
||||
if current_count >= config.requests_per_minute
|
||||
else 0,
|
||||
)
|
||||
|
||||
def reset(self, key: str | None = None) -> None:
|
||||
@@ -148,9 +159,11 @@ class RateLimiter:
|
||||
self.counters.clear()
|
||||
self.configs.clear()
|
||||
|
||||
|
||||
# 全局限流器实例
|
||||
_rate_limiter: RateLimiter | None = None
|
||||
|
||||
|
||||
def get_rate_limiter() -> RateLimiter:
|
||||
"""获取限流器实例"""
|
||||
global _rate_limiter
|
||||
@@ -158,6 +171,7 @@ def get_rate_limiter() -> RateLimiter:
|
||||
_rate_limiter = RateLimiter()
|
||||
return _rate_limiter
|
||||
|
||||
|
||||
# 限流装饰器(用于函数级别限流)
|
||||
def rate_limit(requests_per_minute: int = 60, key_func: Callable | None = None) -> None:
|
||||
"""
|
||||
@@ -178,7 +192,9 @@ def rate_limit(requests_per_minute: int = 60, key_func: Callable | None = None)
|
||||
info = await limiter.is_allowed(key, config)
|
||||
|
||||
if not info.allowed:
|
||||
raise RateLimitExceeded(f"Rate limit exceeded. Try again in {info.retry_after} seconds.")
|
||||
raise RateLimitExceeded(
|
||||
f"Rate limit exceeded. Try again in {info.retry_after} seconds."
|
||||
)
|
||||
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
@@ -189,7 +205,9 @@ def rate_limit(requests_per_minute: int = 60, key_func: Callable | None = None)
|
||||
info = asyncio.run(limiter.is_allowed(key, config))
|
||||
|
||||
if not info.allowed:
|
||||
raise RateLimitExceeded(f"Rate limit exceeded. Try again in {info.retry_after} seconds.")
|
||||
raise RateLimitExceeded(
|
||||
f"Rate limit exceeded. Try again in {info.retry_after} seconds."
|
||||
)
|
||||
|
||||
return func(*args, **kwargs)
|
||||
|
||||
@@ -197,5 +215,6 @@ def rate_limit(requests_per_minute: int = 60, key_func: Callable | None = None)
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
class RateLimitExceeded(Exception):
|
||||
"""限流异常"""
|
||||
|
||||
Reference in New Issue
Block a user