ProjectHentai/nonebot2
6
0
Fork 0
mirror of https://github.com/nonebot/nonebot2.git synced 2025-07-16 02:50:48 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

♿ change websocket client to context manager

Browse Source
This commit is contained in:
yanyongyu
2021-12-26 13:42:13 +08:00
parent 00c2ee8490
commit 7b204d72e6
4 changed files with 37 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
nonebot/drivers/websockets.py
View File

@ -1,4 +1,6 @@
import logging
from typing import AsyncGenerator
from contextlib import asynccontextmanager
from nonebot.typing import overrides
from nonebot.log import LoguruHandler
@ -29,13 +31,15 @@ class Mixin(ForwardMixin):
return await super(Mixin, self).request(setup)
@overrides(ForwardMixin)
async def websocket(self, setup: Request) -> "WebSocket":
ws = await Connect(
@asynccontextmanager
async def websocket(self, setup: Request) -> AsyncGenerator["WebSocket", None]:
connection = Connect(
str(setup.url),
extra_headers=setup.headers.items(),
open_timeout=setup.timeout,
)
return WebSocket(request=setup, websocket=ws)
async with connection as ws:
yield WebSocket(request=setup, websocket=ws)
class WebSocket(BaseWebSocket):