ProjectHentai/nonebot2
6
0
Fork 0
mirror of https://github.com/nonebot/nonebot2.git synced 2025-10-30 22:46:40 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

🐛 fix hook block event receiving bug

Browse Source
This commit is contained in:
yanyongyu
2020-12-28 13:53:24 +08:00
parent c993f15bca
commit 299c259d50
3 changed files with 44 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
nonebot/drivers/__init__.py
View File

@@ -140,7 +140,7 @@ class Driver(abc.ABC):
def on_bot_disconnect(
self,
func: T_WebSocketDisConnectionHook) -> T_WebSocketDisConnectionHook:
func: T_WebSocketDisconnectionHook) -> T_WebSocketDisconnectionHook:
"""
:说明:
@@ -153,30 +153,38 @@ class Driver(abc.ABC):
self._ws_disconnection_hook.add(func)
return func
async def bot_connect(self, bot: "Bot") -> None:
def bot_connect(self, bot: "Bot") -> None:
"""在 WebSocket 连接成功后，调用该函数来注册 bot 对象"""
self._clients[bot.self_id] = bot
coros = list(map(lambda x: x(bot), self._ws_connection_hook))
if coros:
try:
await asyncio.gather(*coros)
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error when running WebSocketConnection hook. "
"Running cancelled!</bg #f8bbd0></r>")
async def bot_disconnect(self, bot: "Bot") -> None:
async def _run_hook(bot: "Bot") -> None:
coros = list(map(lambda x: x(bot), self._ws_connection_hook))
if coros:
try:
await asyncio.gather(*coros)
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error when running WebSocketConnection hook. "
"Running cancelled!</bg #f8bbd0></r>")
asyncio.create_task(_run_hook(bot))
def bot_disconnect(self, bot: "Bot") -> None:
"""在 WebSocket 连接断开后，调用该函数来注销 bot 对象"""
if bot.self_id in self._clients:
del self._clients[bot.self_id]
coros = list(map(lambda x: x(bot), self._ws_disconnection_hook))
if coros:
try:
await asyncio.gather(*coros)
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error when running WebSocketDisConnection hook. "
"Running cancelled!</bg #f8bbd0></r>")
async def _run_hook(bot: "Bot") -> None:
coros = list(map(lambda x: x(bot), self._ws_disconnection_hook))
if coros:
try:
await asyncio.gather(*coros)
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error when running WebSocketDisConnection hook. "
"Running cancelled!</bg #f8bbd0></r>")
asyncio.create_task(_run_hook(bot))
@abc.abstractmethod
def run(self,

5
nonebot/drivers/fastapi.py
View File

@@ -188,11 +188,12 @@ class Driver(BaseDriver):
bot = BotClass(self, "websocket", self.config, x_self_id, websocket=ws)
await ws.accept()
await self.bot_connect(bot)
logger.opt(colors=True).info(
f"WebSocket Connection from <y>{adapter.upper()} "
f"Bot {x_self_id}</y> Accepted!")
self.bot_connect(bot)
try:
while not ws.closed:
data = await ws.receive()
@@ -202,7 +203,7 @@ class Driver(BaseDriver):
asyncio.create_task(bot.handle_message(data))
finally:
await self.bot_disconnect(bot)
self.bot_disconnect(bot)
class WebSocket(BaseWebSocket):