⚗️ add ws connection hook

nonebot/drivers/__init__.py

@@ -6,10 +6,12 @@
 """
 
 import abc
-from typing import Dict, Type, Optional, Callable, TYPE_CHECKING
+import asyncio
+from typing import Set, Dict, Type, Optional, Callable, TYPE_CHECKING
 
 from nonebot.log import logger
 from nonebot.config import Env, Config
+from nonebot.typing import T_WebSocketConnectionHook, T_WebSocketDisconnectionHook
 
 if TYPE_CHECKING:
     from nonebot.adapters import Bot
@@ -25,6 +27,16 @@ class Driver(abc.ABC):
     :类型: ``Dict[str, Type[Bot]]``
     :说明: 已注册的适配器列表
     """
+    _ws_connection_hook: Set[T_WebSocketConnectionHook] = set()
+    """
+    :类型: ``Set[T_WebSocketConnectionHook]``
+    :说明: WebSocket 连接建立时执行的函数
+    """
+    _ws_disconnection_hook: Set[T_WebSocketDisconnectionHook] = set()
+    """
+    :类型: ``Set[T_WebSocketDisconnectionHook]``
+    :说明: WebSocket 连接断开时执行的函数
+    """
 
     @abc.abstractmethod
     def __init__(self, env: Env, config: Config):
@@ -93,8 +105,12 @@ class Driver(abc.ABC):
     @property
     def bots(self) -> Dict[str, "Bot"]:
         """
-        :类型: ``Dict[str, Bot]``
-        :说明: 获取当前所有已连接的 Bot
+        :类型:
+
+          ``Dict[str, Bot]``
+        :说明:
+
+          获取当前所有已连接的 Bot
         """
         return self._clients
 
@@ -108,6 +124,60 @@ class Driver(abc.ABC):
         """注册一个在驱动停止时运行的函数"""
         raise NotImplementedError
 
+    def on_bot_connect(
+            self, func: T_WebSocketConnectionHook) -> T_WebSocketConnectionHook:
+        """
+        :说明:
+
+          装饰一个函数使他在 bot 通过 WebSocket 连接成功时执行。
+
+        :函数参数:
+
+          * ``bot: Bot``: 当前连接上的 Bot 对象
+        """
+        self._ws_connection_hook.add(func)
+        return func
+
+    def on_bot_disconnect(
+            self,
+            func: T_WebSocketDisConnectionHook) -> T_WebSocketDisConnectionHook:
+        """
+        :说明:
+
+          装饰一个函数使他在 bot 通过 WebSocket 连接断开时执行。
+
+        :函数参数:
+
+          * ``bot: Bot``: 当前连接上的 Bot 对象
+        """
+        self._ws_disconnection_hook.add(func)
+        return func
+
+    async 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:
+        """在 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>")
+
     @abc.abstractmethod
     def run(self,
             host: Optional[str] = None,

nonebot/drivers/fastapi.py

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

nonebot/typing.py

@@ -54,6 +54,23 @@ T_StateFactory = Callable[["Bot", "Event"], Awaitable[T_State]]
   事件处理状态 State 类工厂函数
 """
 
+T_WebSocketConnectionHook = Callable[["Bot"], Awaitable[None]]
+"""
+:类型: ``Callable[[Bot], Awaitable[None]]``
+
+:说明:
+
+  WebSocket 连接建立时执行的函数
+"""
+T_WebSocketDisconnectionHook = Callable[["Bot"], Awaitable[None]]
+"""
+:类型: ``Callable[[Bot], Awaitable[None]]``
+
+:说明:
+
+  WebSocket 连接断开时执行的函数
+"""
+
 T_EventPreProcessor = Callable[["Bot", "Event", T_State], Awaitable[None]]
 """
 :类型: ``Callable[[Bot, Event, T_State], Awaitable[None]]``