Feature: 细化内置驱动器请求参数中的超时控制颗粒度 (#3571)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Ju4tCode <42488585+yanyongyu@users.noreply.github.com>
This commit is contained in:
Ailitonia
2025-07-15 10:29:18 +08:00
committed by GitHub
parent bc0682af8f
commit 4ec9bfb7d7
8 changed files with 115 additions and 19 deletions

View File

@ -37,7 +37,14 @@ from nonebot.drivers import (
from nonebot.drivers import WebSocket as BaseWebSocket
from nonebot.drivers.none import Driver as NoneDriver
from nonebot.exception import WebSocketClosed
from nonebot.internal.driver import Cookies, CookieTypes, HeaderTypes, QueryTypes
from nonebot.internal.driver import (
Cookies,
CookieTypes,
HeaderTypes,
QueryTypes,
Timeout,
TimeoutTypes,
)
try:
import aiohttp
@ -56,7 +63,7 @@ class Session(HTTPClientSession):
headers: HeaderTypes = None,
cookies: CookieTypes = None,
version: Union[str, HTTPVersion] = HTTPVersion.H11,
timeout: Optional[float] = None,
timeout: TimeoutTypes = None,
proxy: Optional[str] = None,
):
self._client: Optional[aiohttp.ClientSession] = None
@ -78,7 +85,15 @@ class Session(HTTPClientSession):
else:
raise RuntimeError(f"Unsupported HTTP version: {version}")
self._timeout = timeout
if isinstance(timeout, Timeout):
self._timeout = aiohttp.ClientTimeout(
total=timeout.total,
connect=timeout.connect,
sock_read=timeout.read,
)
else:
self._timeout = aiohttp.ClientTimeout(timeout)
self._proxy = proxy
@property
@ -106,7 +121,14 @@ class Session(HTTPClientSession):
if cookie.value is not None
)
timeout = aiohttp.ClientTimeout(setup.timeout)
if isinstance(setup.timeout, Timeout):
timeout = aiohttp.ClientTimeout(
total=setup.timeout.total,
connect=setup.timeout.connect,
sock_read=setup.timeout.read,
)
else:
timeout = aiohttp.ClientTimeout(setup.timeout)
async with await self.client.request(
setup.method,
@ -149,7 +171,14 @@ class Session(HTTPClientSession):
if cookie.value is not None
)
timeout = aiohttp.ClientTimeout(setup.timeout)
if isinstance(setup.timeout, Timeout):
timeout = aiohttp.ClientTimeout(
total=setup.timeout.total,
connect=setup.timeout.connect,
sock_read=setup.timeout.read,
)
else:
timeout = aiohttp.ClientTimeout(setup.timeout)
async with self.client.request(
setup.method,
@ -226,7 +255,13 @@ class Mixin(HTTPClientMixin, WebSocketClientMixin):
else:
raise RuntimeError(f"Unsupported HTTP version: {setup.version}")
timeout = aiohttp.ClientWSTimeout(ws_close=setup.timeout or 10.0) # type: ignore
if isinstance(setup.timeout, Timeout):
timeout = aiohttp.ClientWSTimeout(
ws_receive=setup.timeout.read, # type: ignore
ws_close=setup.timeout.total, # type: ignore
)
else:
timeout = aiohttp.ClientWSTimeout(ws_close=setup.timeout or 10.0) # type: ignore
async with aiohttp.ClientSession(version=version, trust_env=True) as session:
async with session.ws_connect(
@ -245,7 +280,7 @@ class Mixin(HTTPClientMixin, WebSocketClientMixin):
headers: HeaderTypes = None,
cookies: CookieTypes = None,
version: Union[str, HTTPVersion] = HTTPVersion.H11,
timeout: Optional[float] = None,
timeout: TimeoutTypes = None,
proxy: Optional[str] = None,
) -> Session:
return Session(