diff --git a/nonebot/drivers/websockets.py b/nonebot/drivers/websockets.py index 0645e8ab..28ced115 100644 --- a/nonebot/drivers/websockets.py +++ b/nonebot/drivers/websockets.py @@ -36,7 +36,7 @@ from nonebot.drivers import WebSocket as BaseWebSocket from nonebot.drivers.none import Driver as NoneDriver from nonebot.exception import WebSocketClosed from nonebot.log import LoguruHandler -from nonebot.utils import UNSET, exclude_unset +from nonebot.utils import UNSET, UnsetType, exclude_unset try: from websockets import ClientConnection, ConnectionClosed, connect @@ -77,18 +77,17 @@ class Mixin(WebSocketClientMixin): @override @asynccontextmanager async def websocket(self, setup: Request) -> AsyncGenerator["WebSocket", None]: - timeout_kwargs: dict[str, float | None] = {} + timeout_kwargs: dict[str, float | None | UnsetType] = {} if isinstance(setup.timeout, Timeout): open_timeout = ( setup.timeout.connect or setup.timeout.read or setup.timeout.total ) - timeout_kwargs = exclude_unset( - { - "open_timeout": open_timeout, - "close_timeout": setup.timeout.close, - "ping_timeout": setup.timeout.ping, - } - ) + timeout_kwargs = { + "open_timeout": open_timeout, + "close_timeout": setup.timeout.close, + "ping_timeout": setup.timeout.ping, + } + elif setup.timeout is not UNSET: timeout_kwargs = { "open_timeout": setup.timeout, @@ -99,13 +98,11 @@ class Mixin(WebSocketClientMixin): open_timeout = ( DEFAULT_TIMEOUT.connect or DEFAULT_TIMEOUT.read or DEFAULT_TIMEOUT.total ) - timeout_kwargs = exclude_unset( - { - "open_timeout": open_timeout, - "close_timeout": DEFAULT_TIMEOUT.close, - "ping_timeout": DEFAULT_TIMEOUT.ping, - } - ) + timeout_kwargs = { + "open_timeout": open_timeout, + "close_timeout": DEFAULT_TIMEOUT.close, + "ping_timeout": DEFAULT_TIMEOUT.ping, + } kwargs = exclude_unset( {