🐛 fix hook block event receiving bug

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

View File

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

View File

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

View File

@ -0,0 +1,14 @@
import nonebot
from nonebot.adapters import Bot
driver = nonebot.get_driver()
@driver.on_bot_connect
async def connect(bot: Bot) -> None:
print("Connect", bot)
@driver.on_bot_disconnect
async def disconnect(bot: Bot) -> None:
print("Disconnect", bot)