mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-09-29 00:32:06 +08:00
* ✨ feat(hook): 增强认证钩子和运行时缓存管理 ``` ♻ refactor(hook): 移除未使用的配置项并优化缓存设置 移除 AUTH_HOOKS_CONCURRENCY_LIMIT 配置项,该配置项不再使用 ✨feat(auth_ban): 简化缓存配置并添加实体参数支持 将 BAN_CACHE_TTL 相关配置从动态配置改为常量定义, 移除复杂的 TTL 值转换逻辑,并为 auth_ban 函数添加可选的 entity 参数以支持外部传入实体信息 ♻ refactor(auth_limit): 移除未使用的配置依赖 移除 AUTH_LIMIT_NOTICE_CD 配置项,直接使用常量值 2 作为限制通知冷却时间 📦 依赖更新: update playwright dependency to version 1.57.0 in pyproject.toml and requirements.txt ``` * ✨ feat(auth_checker): 增强插件模块处理和预过滤逻辑,支持用户插件兼容性 * ✨ feat(http_utils): 添加内容缓存机制以优化HTTP响应处理 * ✨ feat(cache): 添加群组插件设置视图缓存类型并更新相关逻辑 * ✨ feat(renderer): 优化渲染引擎,增加内存缓存管理和HTML文档处理逻辑 * ✨ feat(sign_in): 添加好感度排行和好感度总排行命令 * ✨ feat(renderer): 增强浏览器实例管理和模板预处理,支持历史 include 语法兼容 * ✨ feat(renderer): 优化 Playwright 环境检查逻辑,增加结果缓存以提高性能 * ✨ feat(renderer): 增强模板文件渲染策略,优化资产加载路径处理 * 🚨 auto fix by pre-commit hooks * ✨ feat(log): 增加日志内容安全序列化,避免超长 base64 等污染日志 ✨ feat(log_sanitizer): 添加对嵌入超长 base64/data URI 的清理功能 ✨ feat(auth_checker): 添加 Alconna 快捷方式缓存检查,优化路由匹配逻辑 * feat(renderer): 添加渲染结果内存缓存功能以优化性能 feat(help): 实现帮助菜单图像缓存机制 feat(sign_in): 更新HTML卡片生成以支持动画禁用和剪裁 feat(superuser): 在启动时预热超级用户帮助缓存 feat(theme): 优化主题管理器的资源解析缓存 * 🚨 auto fix by pre-commit hooks * feat(renderer): 增加全页面视口最大宽度限制并优化内容尺寸计算 --------- Co-authored-by: ATTomatoo <1126160939@qq.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: HibiKier <45528451+HibiKier@users.noreply.github.com>
794 lines
29 KiB
Python
794 lines
29 KiB
Python
import asyncio
|
||
from collections import OrderedDict
|
||
from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence
|
||
from contextlib import asynccontextmanager
|
||
import os
|
||
from pathlib import Path
|
||
import time
|
||
from typing import Any, ClassVar, cast
|
||
|
||
import aiofiles
|
||
import httpx
|
||
from httpx import AsyncClient, AsyncHTTPTransport, HTTPStatusError, Proxy, Response
|
||
import nonebot
|
||
from rich.progress import (
|
||
BarColumn,
|
||
DownloadColumn,
|
||
Progress,
|
||
TextColumn,
|
||
TransferSpeedColumn,
|
||
)
|
||
import ujson as json
|
||
|
||
from zhenxun.configs.config import BotConfig
|
||
from zhenxun.services.log import logger
|
||
from zhenxun.utils.decorator.retry import Retry
|
||
from zhenxun.utils.exception import AllURIsFailedError
|
||
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
|
||
from zhenxun.utils.user_agent import get_user_agent
|
||
|
||
from .browser import AsyncPlaywright, BrowserIsNone # noqa: F401
|
||
|
||
_SENTINEL = object()
|
||
|
||
driver = nonebot.get_driver()
|
||
_client: AsyncClient | None = None
|
||
|
||
|
||
@PriorityLifecycle.on_startup(priority=0)
|
||
async def _():
|
||
"""
|
||
在Bot启动时初始化全局httpx客户端。
|
||
"""
|
||
global _client
|
||
client_kwargs = {}
|
||
if proxy_url := BotConfig.system_proxy or None:
|
||
try:
|
||
version_parts = httpx.__version__.split(".")
|
||
major = int("".join(c for c in version_parts[0] if c.isdigit()))
|
||
minor = (
|
||
int("".join(c for c in version_parts[1] if c.isdigit()))
|
||
if len(version_parts) > 1
|
||
else 0
|
||
)
|
||
if (major, minor) >= (0, 28):
|
||
client_kwargs["proxy"] = proxy_url
|
||
else:
|
||
client_kwargs["proxies"] = proxy_url
|
||
except (ValueError, IndexError):
|
||
client_kwargs["proxy"] = proxy_url
|
||
logger.warning(
|
||
f"无法解析 httpx 版本 '{httpx.__version__}',"
|
||
"将默认使用新版 'proxy' 参数语法。"
|
||
)
|
||
|
||
_client = get_async_client(
|
||
headers=get_user_agent(),
|
||
follow_redirects=True,
|
||
**client_kwargs,
|
||
)
|
||
|
||
logger.info("全局 httpx.AsyncClient 已启动。", "HTTPClient")
|
||
|
||
|
||
@driver.on_shutdown
|
||
async def _():
|
||
"""
|
||
在Bot关闭时关闭全局httpx客户端。
|
||
"""
|
||
if _client:
|
||
await _client.aclose()
|
||
logger.info("全局 httpx.AsyncClient 已关闭。", "HTTPClient")
|
||
|
||
|
||
def get_client() -> AsyncClient:
|
||
"""
|
||
获取全局 httpx.AsyncClient 实例。
|
||
"""
|
||
global _client
|
||
if not _client:
|
||
if not os.environ.get("PYTEST_CURRENT_TEST"):
|
||
raise RuntimeError("全局 httpx.AsyncClient 未初始化,请检查启动流程。")
|
||
# 在测试环境中创建临时客户端
|
||
logger.warning("在测试环境中创建临时HTTP客户端", "HTTPClient")
|
||
_client = get_async_client(
|
||
headers=get_user_agent(),
|
||
follow_redirects=True,
|
||
)
|
||
return _client
|
||
|
||
|
||
def get_async_client(
|
||
proxies: dict[str, str] | str | None = None,
|
||
proxy: str | None = None,
|
||
verify: bool = False,
|
||
**kwargs,
|
||
) -> httpx.AsyncClient:
|
||
"""
|
||
[向后兼容] 创建 httpx.AsyncClient 实例的工厂函数。
|
||
此函数完全保留了旧版本的接口,确保现有代码无需修改即可使用。
|
||
"""
|
||
transport = kwargs.pop("transport", None) or AsyncHTTPTransport(verify=verify)
|
||
if proxies:
|
||
if isinstance(proxies, str):
|
||
proxies = {"http://": proxies, "https://": proxies}
|
||
http_proxy = proxies.get("http://")
|
||
https_proxy = proxies.get("https://")
|
||
return httpx.AsyncClient(
|
||
mounts={
|
||
"http://": AsyncHTTPTransport(
|
||
proxy=Proxy(http_proxy) if http_proxy else None
|
||
),
|
||
"https://": AsyncHTTPTransport(
|
||
proxy=Proxy(https_proxy) if https_proxy else None
|
||
),
|
||
},
|
||
transport=transport,
|
||
**kwargs,
|
||
)
|
||
elif proxy:
|
||
return httpx.AsyncClient(
|
||
mounts={
|
||
"http://": AsyncHTTPTransport(proxy=Proxy(proxy)),
|
||
"https://": AsyncHTTPTransport(proxy=Proxy(proxy)),
|
||
},
|
||
transport=transport,
|
||
**kwargs,
|
||
)
|
||
return httpx.AsyncClient(transport=transport, **kwargs)
|
||
|
||
|
||
class AsyncHttpx:
|
||
"""
|
||
高性能异步HTTP客户端工具类。
|
||
|
||
特性:
|
||
- 全局共享连接池,提升性能
|
||
- 支持临时客户端配置(代理、超时等)
|
||
- 内置重试机制和多URL回退
|
||
- 提供JSON解析和文件下载功能
|
||
"""
|
||
|
||
CLIENT_KEY: ClassVar[list[str]] = [
|
||
"use_proxy",
|
||
"proxies",
|
||
"proxy",
|
||
"verify",
|
||
]
|
||
|
||
default_proxy: ClassVar[dict[str, str] | None] = (
|
||
{
|
||
"http://": BotConfig.system_proxy,
|
||
"https://": BotConfig.system_proxy,
|
||
}
|
||
if BotConfig.system_proxy
|
||
else None
|
||
)
|
||
|
||
_CONTENT_CACHE_TTL: ClassVar[float] = 3.0
|
||
_CONTENT_CACHE_MAX_ITEMS: ClassVar[int] = 256
|
||
_CONTENT_CACHE_MAX_BYTES: ClassVar[int] = 2 * 1024 * 1024
|
||
_content_cache: ClassVar[OrderedDict[str, tuple[float, bytes]]] = OrderedDict()
|
||
_content_inflight: ClassVar[dict[str, asyncio.Task[Response]]] = {}
|
||
_content_cache_lock: ClassVar[asyncio.Lock] = asyncio.Lock()
|
||
|
||
@classmethod
|
||
def _is_probably_image_url(cls, url: str) -> bool:
|
||
lower_url = url.lower()
|
||
if any(
|
||
ext in lower_url
|
||
for ext in (
|
||
".jpg",
|
||
".jpeg",
|
||
".png",
|
||
".webp",
|
||
".gif",
|
||
".bmp",
|
||
".avif",
|
||
".heic",
|
||
)
|
||
):
|
||
return True
|
||
return "qpic.cn" in lower_url or "qlogo.cn" in lower_url
|
||
|
||
@classmethod
|
||
def _get_cached_content_nolock(cls, key: str) -> bytes | None:
|
||
entry = cls._content_cache.get(key)
|
||
if not entry:
|
||
return None
|
||
expire_at, content = entry
|
||
if expire_at <= time.monotonic():
|
||
cls._content_cache.pop(key, None)
|
||
return None
|
||
cls._content_cache.move_to_end(key)
|
||
return content
|
||
|
||
@classmethod
|
||
def _cleanup_content_cache_nolock(cls) -> None:
|
||
now = time.monotonic()
|
||
while cls._content_cache:
|
||
expire_at, _ = next(iter(cls._content_cache.values()))
|
||
if expire_at > now:
|
||
break
|
||
cls._content_cache.popitem(last=False)
|
||
while len(cls._content_cache) > cls._CONTENT_CACHE_MAX_ITEMS:
|
||
cls._content_cache.popitem(last=False)
|
||
|
||
@classmethod
|
||
async def _try_cache_content(cls, key: str, response: Response) -> None:
|
||
content = response.content
|
||
if not content or len(content) > cls._CONTENT_CACHE_MAX_BYTES:
|
||
return
|
||
content_type = response.headers.get("content-type", "").lower()
|
||
is_image = content_type.startswith("image/") or cls._is_probably_image_url(key)
|
||
if not is_image:
|
||
return
|
||
async with cls._content_cache_lock:
|
||
cls._content_cache[key] = (
|
||
time.monotonic() + cls._CONTENT_CACHE_TTL,
|
||
content,
|
||
)
|
||
cls._content_cache.move_to_end(key)
|
||
cls._cleanup_content_cache_nolock()
|
||
|
||
@classmethod
|
||
def _prepare_temporary_client_config(cls, client_kwargs: dict) -> dict:
|
||
"""
|
||
[向后兼容] 处理旧式的客户端kwargs,将其转换为get_async_client可用的配置。
|
||
主要负责处理 use_proxy 标志,这是为了兼容旧版本代码中使用的 use_proxy 参数。
|
||
"""
|
||
final_config = client_kwargs.copy()
|
||
|
||
use_proxy = final_config.pop("use_proxy", True)
|
||
|
||
if "proxies" not in final_config and "proxy" not in final_config:
|
||
final_config["proxies"] = cls.default_proxy if use_proxy else None
|
||
return final_config
|
||
|
||
@classmethod
|
||
def _split_kwargs(cls, kwargs: dict) -> tuple[dict, dict]:
|
||
"""[优化] 分离客户端配置和请求参数,使逻辑更清晰。"""
|
||
client_kwargs = {k: v for k, v in kwargs.items() if k in cls.CLIENT_KEY}
|
||
request_kwargs = {k: v for k, v in kwargs.items() if k not in cls.CLIENT_KEY}
|
||
return client_kwargs, request_kwargs
|
||
|
||
@classmethod
|
||
@asynccontextmanager
|
||
async def _get_active_client_context(
|
||
cls, client: AsyncClient | None = None, **kwargs
|
||
) -> AsyncGenerator[AsyncClient, None]:
|
||
"""
|
||
内部辅助方法,根据 kwargs 决定并提供一个活动的 HTTP 客户端。
|
||
- 如果 kwargs 中有客户端配置,则创建并返回一个临时客户端。
|
||
- 否则,返回传入的 client 或全局客户端。
|
||
- 自动处理临时客户端的关闭。
|
||
"""
|
||
if kwargs:
|
||
logger.debug(f"为单次请求创建临时客户端,配置: {kwargs}")
|
||
temp_client_config = cls._prepare_temporary_client_config(kwargs)
|
||
async with get_async_client(**temp_client_config) as temp_client:
|
||
yield temp_client
|
||
else:
|
||
yield client or get_client()
|
||
|
||
@Retry.simple(log_name="内部HTTP请求")
|
||
async def _execute_request_inner(
|
||
self, client: AsyncClient, method: str, url: str, **kwargs
|
||
) -> Response:
|
||
"""
|
||
[内部] 执行单次HTTP请求的私有核心方法,被重试装饰器包裹。
|
||
"""
|
||
return await client.request(method, url, **kwargs)
|
||
|
||
@classmethod
|
||
async def _single_request(
|
||
cls, method: str, url: str, *, client: AsyncClient | None = None, **kwargs
|
||
) -> Response:
|
||
"""
|
||
执行单次HTTP请求的私有方法,内置了默认的重试逻辑。
|
||
accept_status_codes: 若提供,这些状态码不视为错误,不调用 raise_for_status。
|
||
"""
|
||
client_kwargs, request_kwargs = cls._split_kwargs(kwargs)
|
||
accept_status_codes: Sequence[int] | None = request_kwargs.pop(
|
||
"accept_status_codes", None
|
||
)
|
||
|
||
async with cls._get_active_client_context(
|
||
client=client, **client_kwargs
|
||
) as active_client:
|
||
response = await cls()._execute_request_inner(
|
||
active_client, method, url, **request_kwargs
|
||
)
|
||
if (
|
||
accept_status_codes is None
|
||
or response.status_code not in accept_status_codes
|
||
):
|
||
response.raise_for_status()
|
||
return response
|
||
|
||
@classmethod
|
||
async def _execute_with_fallbacks(
|
||
cls,
|
||
urls: str | list[str],
|
||
worker: Callable[..., Awaitable[Any]],
|
||
*,
|
||
client: AsyncClient | None = None,
|
||
**kwargs,
|
||
) -> Any:
|
||
"""
|
||
通用执行器,按顺序尝试多个URL,直到成功。
|
||
|
||
参数:
|
||
urls: 单个URL或URL列表。
|
||
worker: 一个接受单个URL和其他kwargs并执行请求的协程函数。
|
||
client: 可选的HTTP客户端。
|
||
**kwargs: 传递给worker的额外参数。
|
||
"""
|
||
url_list = [urls] if isinstance(urls, str) else urls
|
||
exceptions = []
|
||
|
||
for i, url in enumerate(url_list):
|
||
try:
|
||
result = await worker(url, client=client, **kwargs)
|
||
if i > 0:
|
||
logger.info(
|
||
f"成功从镜像 '{url}' 获取资源 "
|
||
f"(在尝试了 {i} 个失败的镜像之后)。",
|
||
"AsyncHttpx:FallbackExecutor",
|
||
)
|
||
return result
|
||
except Exception as e:
|
||
if isinstance(e, HTTPStatusError):
|
||
status = getattr(e.response, "status_code", "?")
|
||
try:
|
||
body_text = getattr(e.response, "text", None)
|
||
if body_text is not None and len(body_text) > 2000:
|
||
body_text = body_text[:2000] + "...(truncated)"
|
||
except Exception:
|
||
body_text = "<unavailable>"
|
||
logger.debug(
|
||
f"请求失败: {url} {status} {body_text}",
|
||
"AsyncHttpx:FallbackExecutor",
|
||
)
|
||
exceptions.append(e)
|
||
if url != url_list[-1]:
|
||
logger.warning(
|
||
f"Worker '{worker.__name__}' on {url} failed, trying next. "
|
||
f"Error: {e.__class__.__name__}",
|
||
"AsyncHttpx:FallbackExecutor",
|
||
)
|
||
|
||
raise AllURIsFailedError(url_list, exceptions)
|
||
|
||
@classmethod
|
||
async def get(
|
||
cls,
|
||
url: str | list[str],
|
||
*,
|
||
follow_redirects: bool = True,
|
||
check_status_code: int | None = None,
|
||
accept_status_codes: Sequence[int] | None = None,
|
||
client: AsyncClient | None = None,
|
||
**kwargs,
|
||
) -> Response:
|
||
"""发送 GET 请求,并返回第一个成功的响应。
|
||
|
||
参数:
|
||
url: 单个请求 URL 或一个 URL 列表。
|
||
follow_redirects: 是否跟随重定向。
|
||
check_status_code: (可选) 若提供,将检查响应状态码是否匹配,否则抛出异常。
|
||
accept_status_codes: (可选) 这些状态码不视为错误,
|
||
如 (302,) 用于接受重定向响应。
|
||
client: (可选) 指定一个活动的HTTP客户端实例。若提供,则忽略
|
||
`**kwargs`中的客户端配置。
|
||
**kwargs: 其他所有传递给 httpx.get 的参数 (如 `params`, `headers`,
|
||
`timeout`)。如果包含 `proxies`, `verify` 等客户端配置参数,
|
||
将创建一个临时客户端。
|
||
|
||
返回:
|
||
Response: httpx 的响应对象。
|
||
|
||
异常:
|
||
AllURIsFailedError: 当所有提供的URL都请求失败时抛出。
|
||
"""
|
||
|
||
async def worker(current_url: str, **worker_kwargs) -> Response:
|
||
logger.info(f"开始获取 {current_url}..", "AsyncHttpx:get")
|
||
if accept_status_codes is not None:
|
||
worker_kwargs["accept_status_codes"] = accept_status_codes
|
||
response = await cls._single_request(
|
||
"GET", current_url, follow_redirects=follow_redirects, **worker_kwargs
|
||
)
|
||
if check_status_code and response.status_code != check_status_code:
|
||
raise HTTPStatusError(
|
||
f"状态码错误: {response.status_code}!={check_status_code}",
|
||
request=response.request,
|
||
response=response,
|
||
)
|
||
return response
|
||
|
||
return await cls._execute_with_fallbacks(url, worker, client=client, **kwargs)
|
||
|
||
@classmethod
|
||
async def head(
|
||
cls, url: str | list[str], *, client: AsyncClient | None = None, **kwargs
|
||
) -> Response:
|
||
"""发送 HEAD 请求,并返回第一个成功的响应。"""
|
||
|
||
async def worker(current_url: str, **worker_kwargs) -> Response:
|
||
return await cls._single_request("HEAD", current_url, **worker_kwargs)
|
||
|
||
return await cls._execute_with_fallbacks(url, worker, client=client, **kwargs)
|
||
|
||
@classmethod
|
||
async def post(
|
||
cls,
|
||
url: str | list[str],
|
||
*,
|
||
accept_status_codes: Sequence[int] | None = None,
|
||
client: AsyncClient | None = None,
|
||
**kwargs,
|
||
) -> Response:
|
||
"""发送 POST 请求,并返回第一个成功的响应。
|
||
accept_status_codes: (可选) 这些状态码不视为错误,如 (302,) 用于接受重定向响应。
|
||
"""
|
||
|
||
async def worker(current_url: str, **worker_kwargs) -> Response:
|
||
if accept_status_codes is not None:
|
||
worker_kwargs["accept_status_codes"] = accept_status_codes
|
||
return await cls._single_request("POST", current_url, **worker_kwargs)
|
||
|
||
return await cls._execute_with_fallbacks(url, worker, client=client, **kwargs)
|
||
|
||
@classmethod
|
||
async def get_content(
|
||
cls, url: str | list[str], *, client: AsyncClient | None = None, **kwargs
|
||
) -> bytes:
|
||
"""获取指定 URL 的二进制内容。"""
|
||
if not isinstance(url, str):
|
||
res = await cls.get(url, client=client, **kwargs)
|
||
return res.content
|
||
|
||
cache_key = url
|
||
async with cls._content_cache_lock:
|
||
cached = cls._get_cached_content_nolock(cache_key)
|
||
if cached is not None:
|
||
return cached
|
||
task = cls._content_inflight.get(cache_key)
|
||
if task is None:
|
||
task = asyncio.create_task(cls.get(url, client=client, **kwargs))
|
||
cls._content_inflight[cache_key] = task
|
||
|
||
def _cleanup_inflight(
|
||
_: asyncio.Task[Response], key: str = cache_key
|
||
) -> None:
|
||
cls._content_inflight.pop(key, None)
|
||
|
||
task.add_done_callback(_cleanup_inflight)
|
||
|
||
res = await task
|
||
await cls._try_cache_content(cache_key, res)
|
||
return res.content
|
||
|
||
@classmethod
|
||
@Retry.api(
|
||
log_name="JSON请求",
|
||
exception=(json.JSONDecodeError,),
|
||
return_on_failure=_SENTINEL,
|
||
)
|
||
async def _request_and_parse_json(
|
||
cls, method: str, url: str, *, client: AsyncClient | None = None, **kwargs
|
||
) -> Any:
|
||
"""
|
||
[私有] 执行单个HTTP请求并解析JSON,用于内部统一处理。
|
||
"""
|
||
client_kwargs, request_kwargs = cls._split_kwargs(kwargs)
|
||
|
||
async with cls._get_active_client_context(
|
||
client=client, **client_kwargs
|
||
) as active_client:
|
||
response = await active_client.request(method, url, **request_kwargs)
|
||
response.raise_for_status()
|
||
return response.json()
|
||
|
||
@classmethod
|
||
async def get_json(
|
||
cls,
|
||
url: str | list[str],
|
||
*,
|
||
default: Any = None,
|
||
raise_on_failure: bool = False,
|
||
client: AsyncClient | None = None,
|
||
**kwargs,
|
||
) -> Any:
|
||
"""
|
||
发送GET请求并自动解析为JSON,支持重试和多链接尝试。
|
||
|
||
参数:
|
||
url: 单个请求 URL 或一个备用 URL 列表。
|
||
default: (可选) 当所有尝试都失败时返回的默认值,默认为None。
|
||
raise_on_failure: (可选) 如果为 True, 当所有尝试失败时将抛出
|
||
`AllURIsFailedError` 异常, 默认为 False.
|
||
client: (可选) 指定的HTTP客户端。
|
||
**kwargs: 其他所有传递给 httpx.get 的参数。
|
||
例如 `params`, `headers`, `timeout`等。
|
||
|
||
返回:
|
||
Any: 解析后的JSON数据,或在失败时返回 `default` 值。
|
||
|
||
异常:
|
||
AllURIsFailedError: 当 `raise_on_failure` 为 True 且所有URL都请求失败时抛出
|
||
"""
|
||
|
||
async def worker(current_url: str, **worker_kwargs):
|
||
logger.debug(f"开始GET JSON: {current_url}", "AsyncHttpx:get_json")
|
||
return await cls._request_and_parse_json(
|
||
"GET", current_url, **worker_kwargs
|
||
)
|
||
|
||
try:
|
||
result = await cls._execute_with_fallbacks(
|
||
url, worker, client=client, **kwargs
|
||
)
|
||
return default if result is _SENTINEL else result
|
||
except AllURIsFailedError as e:
|
||
logger.error(f"所有URL的JSON GET均失败: {e}", "AsyncHttpx:get_json")
|
||
if raise_on_failure:
|
||
raise e
|
||
return default
|
||
|
||
@classmethod
|
||
async def post_json(
|
||
cls,
|
||
url: str | list[str],
|
||
*,
|
||
json: Any = None,
|
||
data: Any = None,
|
||
default: Any = None,
|
||
raise_on_failure: bool = False,
|
||
client: AsyncClient | None = None,
|
||
**kwargs,
|
||
) -> Any:
|
||
"""
|
||
发送POST请求并自动解析为JSON,功能与 get_json 类似。
|
||
|
||
参数:
|
||
url: 单个请求 URL 或一个备用 URL 列表。
|
||
json: (可选) 作为请求体发送的JSON数据。
|
||
data: (可选) 作为请求体发送的表单数据。
|
||
default: (可选) 当所有尝试都失败时返回的默认值,默认为None。
|
||
raise_on_failure: (可选) 如果为 True, 当所有尝试失败时将抛出
|
||
AllURIsFailedError 异常, 默认为 False.
|
||
client: (可选) 指定的HTTP客户端。
|
||
**kwargs: 其他所有传递给 httpx.post 的参数。
|
||
|
||
返回:
|
||
Any: 解析后的JSON数据,或在失败时返回 `default` 值。
|
||
"""
|
||
if json is not None:
|
||
kwargs["json"] = json
|
||
if data is not None:
|
||
kwargs["data"] = data
|
||
|
||
async def worker(current_url: str, **worker_kwargs):
|
||
logger.debug(f"开始POST JSON: {current_url}", "AsyncHttpx:post_json")
|
||
return await cls._request_and_parse_json(
|
||
"POST", current_url, **worker_kwargs
|
||
)
|
||
|
||
try:
|
||
result = await cls._execute_with_fallbacks(
|
||
url, worker, client=client, **kwargs
|
||
)
|
||
return default if result is _SENTINEL else result
|
||
except AllURIsFailedError as e:
|
||
logger.error(f"所有URL的JSON POST均失败: {e}", "AsyncHttpx:post_json")
|
||
if raise_on_failure:
|
||
raise e
|
||
return default
|
||
|
||
@classmethod
|
||
@Retry.api(log_name="文件下载(流式)")
|
||
async def _stream_download(
|
||
cls, url: str, path: Path, *, client: AsyncClient | None = None, **kwargs
|
||
) -> None:
|
||
"""
|
||
执行单个流式下载的私有方法,被重试装饰器包裹。
|
||
"""
|
||
client_kwargs, request_kwargs = cls._split_kwargs(kwargs)
|
||
show_progress = request_kwargs.pop("show_progress", False)
|
||
|
||
async with cls._get_active_client_context(
|
||
client=client, **client_kwargs
|
||
) as active_client:
|
||
async with active_client.stream("GET", url, **request_kwargs) as response:
|
||
response.raise_for_status()
|
||
total = int(response.headers.get("Content-Length", 0))
|
||
|
||
if show_progress:
|
||
with Progress(
|
||
TextColumn(path.name),
|
||
"[progress.percentage]{task.percentage:>3.0f}%",
|
||
BarColumn(bar_width=None),
|
||
DownloadColumn(),
|
||
TransferSpeedColumn(),
|
||
) as progress:
|
||
task_id = progress.add_task("Download", total=total)
|
||
async with aiofiles.open(path, "wb") as f:
|
||
async for chunk in response.aiter_bytes():
|
||
await f.write(chunk)
|
||
progress.update(task_id, advance=len(chunk))
|
||
else:
|
||
async with aiofiles.open(path, "wb") as f:
|
||
async for chunk in response.aiter_bytes():
|
||
await f.write(chunk)
|
||
|
||
@classmethod
|
||
async def download_file(
|
||
cls,
|
||
url: str | list[str],
|
||
path: str | Path,
|
||
*,
|
||
stream: bool = False,
|
||
show_progress: bool = False,
|
||
client: AsyncClient | None = None,
|
||
**kwargs,
|
||
) -> bool:
|
||
"""下载文件到指定路径。
|
||
|
||
说明:
|
||
支持多链接尝试和流式下载(带进度条)。
|
||
|
||
参数:
|
||
url: 单个文件 URL 或一个备用 URL 列表。
|
||
path: 文件保存的本地路径。
|
||
stream: (可选) 是否使用流式下载,适用于大文件,默认为 False。
|
||
show_progress: (可选) 当 stream=True 时,是否显示下载进度条。默认为 False。
|
||
client: (可选) 指定的HTTP客户端。
|
||
**kwargs: 其他所有传递给 get() 方法或 httpx.stream() 的参数。
|
||
|
||
返回:
|
||
bool: 是否下载成功。
|
||
"""
|
||
path = Path(path)
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
async def worker(current_url: str, **worker_kwargs) -> bool:
|
||
if not stream:
|
||
content = await cls.get_content(current_url, **worker_kwargs)
|
||
async with aiofiles.open(path, "wb") as f:
|
||
await f.write(content)
|
||
else:
|
||
await cls._stream_download(
|
||
current_url, path, show_progress=show_progress, **worker_kwargs
|
||
)
|
||
|
||
logger.info(
|
||
f"下载 {current_url} 成功 -> {path.absolute()}",
|
||
"AsyncHttpx:download",
|
||
)
|
||
return True
|
||
|
||
try:
|
||
return await cls._execute_with_fallbacks(
|
||
url, worker, client=client, **kwargs
|
||
)
|
||
except AllURIsFailedError:
|
||
logger.error(
|
||
f"所有URL下载均失败 -> {path.absolute()}", "AsyncHttpx:download"
|
||
)
|
||
return False
|
||
|
||
@classmethod
|
||
async def gather_download_file(
|
||
cls,
|
||
url_list: Sequence[list[str] | str],
|
||
path_list: Sequence[str | Path],
|
||
*,
|
||
limit_async_number: int = 5,
|
||
**kwargs,
|
||
) -> list[bool]:
|
||
"""并发下载多个文件,支持为每个文件提供备用镜像链接。
|
||
|
||
参数:
|
||
url_list: 包含所有文件下载任务的列表。每个元素可以是:
|
||
- 一个字符串 (str): 代表该任务的唯一URL。
|
||
- 一个字符串列表 (list[str]): 代表该任务的多个备用/镜像URL。
|
||
path_list: 与 url_list 对应的文件保存路径列表。
|
||
limit_async_number: (可选) 最大并发下载数,默认为 5。
|
||
**kwargs: 其他所有传递给 download_file() 方法的参数。
|
||
|
||
返回:
|
||
list[bool]: 对应每个下载任务是否成功。
|
||
"""
|
||
if len(url_list) != len(path_list):
|
||
raise ValueError("URL 列表和路径列表的长度必须相等")
|
||
|
||
semaphore = asyncio.Semaphore(limit_async_number)
|
||
|
||
async def _download_with_semaphore(
|
||
urls_for_one_path: str | list[str], path: str | Path
|
||
):
|
||
async with semaphore:
|
||
return await cls.download_file(urls_for_one_path, path, **kwargs)
|
||
|
||
tasks = [
|
||
_download_with_semaphore(url_group, path)
|
||
for url_group, path in zip(url_list, path_list)
|
||
]
|
||
|
||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||
|
||
final_results = []
|
||
for i, result in enumerate(results):
|
||
if isinstance(result, Exception):
|
||
url_info = (
|
||
url_list[i]
|
||
if isinstance(url_list[i], str)
|
||
else ", ".join(url_list[i])
|
||
)
|
||
logger.error(f"并发下载任务 ({url_info}) 时发生错误", e=result)
|
||
final_results.append(False)
|
||
else:
|
||
final_results.append(cast(bool, result))
|
||
|
||
return final_results
|
||
|
||
@classmethod
|
||
async def get_fastest_mirror(cls, url_list: list[str]) -> list[str]:
|
||
"""测试并返回最快的镜像地址。
|
||
|
||
参数:
|
||
url_list: 需要测试的镜像 URL 列表。
|
||
|
||
返回:
|
||
list[str]: 按从快到慢的顺序包含了所有可用的 URL。
|
||
"""
|
||
assert url_list
|
||
|
||
async def head_mirror(client: type[AsyncHttpx], url: str) -> dict[str, Any]:
|
||
begin_time = time.time()
|
||
|
||
response = await client.head(url=url, timeout=6)
|
||
|
||
elapsed_time = (time.time() - begin_time) * 1000
|
||
content_length = int(response.headers.get("content-length", 0))
|
||
|
||
return {
|
||
"url": url,
|
||
"elapsed_time": elapsed_time,
|
||
"content_length": content_length,
|
||
}
|
||
|
||
logger.debug(f"开始获取最快镜像,可能需要一段时间... | URL列表:{url_list}")
|
||
results = await asyncio.gather(
|
||
*(head_mirror(cls, url) for url in url_list),
|
||
return_exceptions=True,
|
||
)
|
||
_results: list[dict[str, Any]] = []
|
||
for result in results:
|
||
if isinstance(result, BaseException):
|
||
logger.warning(f"获取镜像失败,错误:{result}")
|
||
else:
|
||
logger.debug(f"获取镜像成功,结果:{result}")
|
||
_results.append(result)
|
||
_results = sorted(iter(_results), key=lambda r: r["elapsed_time"])
|
||
return [result["url"] for result in _results]
|
||
|
||
@classmethod
|
||
@asynccontextmanager
|
||
async def temporary_client(cls, **kwargs) -> AsyncGenerator[AsyncClient, None]:
|
||
"""
|
||
创建一个临时的、可配置的HTTP客户端上下文,并直接返回该客户端实例。
|
||
|
||
参数:
|
||
**kwargs: 所有传递给 `httpx.AsyncClient` 构造函数的参数。
|
||
例如: `proxies`, `headers`, `verify`, `timeout`,
|
||
`follow_redirects`。
|
||
|
||
返回:
|
||
httpx.AsyncClient: 一个配置好的、临时的客户端实例。
|
||
"""
|
||
async with get_async_client(**kwargs) as client:
|
||
yield client
|