ProjectHentai/nonebot2
6
0
Fork 0
mirror of https://github.com/nonebot/nonebot2.git synced 2025-07-27 16:21:28 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

🐛 Fix: 修复 bot hook 缺少依赖缓存和上下文管理 (#1826)

Browse Source
This commit is contained in:
Ju4tCode
2023-03-20 22:37:57 +08:00
committed by GitHub
parent 05a6af46b9
commit 78bbf9e623
2 changed files with 106 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
nonebot/internal/driver/driver.py
View File

@ -1,6 +1,6 @@
import abc
import asyncio
from contextlib import asynccontextmanager
from contextlib import AsyncExitStack, asynccontextmanager
from typing import TYPE_CHECKING, Any, Set, Dict, Type, Callable, AsyncGenerator
from nonebot.log import logger
@ -8,8 +8,12 @@ from nonebot.config import Env, Config
from nonebot.dependencies import Dependent
from nonebot.exception import SkippedException
from nonebot.utils import escape_tag, run_coro_with_catch
from nonebot.typing import T_BotConnectionHook, T_BotDisconnectionHook
from nonebot.internal.params import BotParam, DependParam, DefaultParam
from nonebot.typing import (
T_DependencyCache,
T_BotConnectionHook,
T_BotDisconnectionHook,
)
from .model import Request, Response, WebSocket, HTTPServerSetup, WebSocketServerSetup
@ -135,20 +139,22 @@ class Driver(abc.ABC):
self._bots[bot.self_id] = bot
async def _run_hook(bot: "Bot") -> None:
coros = list(
map(
lambda x: run_coro_with_catch(x(bot=bot), (SkippedException,)),
self._bot_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>"
dependency_cache: T_DependencyCache = {}
async with AsyncExitStack() as stack:
if coros := [
run_coro_with_catch(
hook(bot=bot, stack=stack, dependency_cache=dependency_cache),
(SkippedException,),
)
for hook in self._bot_connection_hook
]:
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))
@ -158,20 +164,22 @@ class Driver(abc.ABC):
del self._bots[bot.self_id]
async def _run_hook(bot: "Bot") -> None:
coros = list(
map(
lambda x: run_coro_with_catch(x(bot=bot), (SkippedException,)),
self._bot_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>"
dependency_cache: T_DependencyCache = {}
async with AsyncExitStack() as stack:
if coros := [
run_coro_with_catch(
hook(bot=bot, stack=stack, dependency_cache=dependency_cache),
(SkippedException,),
)
for hook in self._bot_disconnection_hook
]:
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))