mirror of
https://github.com/nonebot/nonebot2.git
synced 2024-11-28 03:05:25 +08:00
🐛 fix hook block event receiving bug
This commit is contained in:
parent
c993f15bca
commit
299c259d50
@ -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,
|
||||||
|
@ -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):
|
||||||
|
14
tests/test_plugins/test_ws_hook.py
Normal file
14
tests/test_plugins/test_ws_hook.py
Normal 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)
|
Loading…
Reference in New Issue
Block a user