ProjectHentai/nonebot2
6
0
Fork 0
mirror of https://github.com/nonebot/nonebot2.git synced 2024-11-30 17:15:08 +08: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:45:26 +08:00
parent 7b204d72e6
commit e64f399370
1 changed files with 6 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
nonebot/adapters/_adapter.py
View File

@ -1,5 +1,6 @@
import abc
from typing import Any, Dict
from contextlib import asynccontextmanager
from typing import Any, Dict, AsyncGenerator
from ._bot import Bot
from nonebot.config import Config
@ -52,10 +53,12 @@ class Adapter(abc.ABC):
raise TypeError("Current driver does not support http client")
return await self.driver.request(setup)
async def websocket(self, setup: Request) -> WebSocket:
@asynccontextmanager
async def websocket(self, setup: Request) -> AsyncGenerator[WebSocket, None]:
if not isinstance(self.driver, ForwardDriver):
raise TypeError("Current driver does not support websocket client")
return await self.driver.websocket(setup)
async with self.driver.websocket(setup) as ws:
yield ws
@abc.abstractmethod
async def _call_api(self, bot: Bot, api: str, **data) -> Any: