ProjectHentai/nonebot2
6
0
Fork 0
mirror of https://github.com/nonebot/nonebot2.git synced 2025-07-15 10:31:04 +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

9
nonebot/drivers/httpx.py
View File

@ -1,3 +1,6 @@
from typing import AsyncGenerator
from contextlib import asynccontextmanager
from nonebot.typing import overrides
from nonebot.drivers._block_driver import BlockDriver
from nonebot.drivers import (
@ -48,8 +51,10 @@ class Mixin(ForwardMixin):
)
@overrides(ForwardMixin)
async def websocket(self, setup: Request) -> WebSocket:
return await super(Mixin, self).websocket(setup)
@asynccontextmanager
async def websocket(self, setup: Request) -> AsyncGenerator[WebSocket, None]:
async with super(Mixin, self).websocket(setup) as ws:
yield ws
Driver = combine_driver(BlockDriver, Mixin)