change websocket client to context manager

This commit is contained in:
yanyongyu 2021-12-26 13:45:26 +08:00
parent 7b204d72e6
commit e64f399370

View File

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