✨ Extension.inject

Author: RF-Tar-Railt

src/nonebot_plugin_alconna/extension.py

@@ -2,24 +2,30 @@
 
 from abc import ABCMeta, abstractmethod
 import asyncio
+from contextlib import AsyncExitStack
 from dataclasses import dataclass
 import functools
 import importlib as imp
+import inspect
 import re
-from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypeVar, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypeVar, Union, final, overload
 from weakref import finalize
 
 from arclet.alconna import Alconna, Arparma
 from nonebot import get_plugin_config
 from nonebot.adapters import Bot, Event, Message
 from nonebot.compat import PydanticUndefined
-from nonebot.typing import T_State
+from nonebot.dependencies import Dependent, Param
+from nonebot.internal.params import DependencyCache, DependParam, DependsInner
+from nonebot.typing import T_State, _DependentCallable
+from pydantic.fields import FieldInfo
 from tarina import LRU, lang
 
 from .config import Config
 from .uniseg import UniMessage, get_message_id
 
 OutputType = Literal["help", "shortcut", "completion", "error"]
+T = TypeVar("T")
 TM = TypeVar("TM", bound=Union[str, Message, UniMessage])
 TE = TypeVar("TE", bound=Event)
 
@@ -57,6 +63,8 @@ def __init_subclass__(cls, **kwargs):
             "catch": cls.catch != Extension.catch and cls.before_catch != Extension.before_catch,
         }
 
+    executor: ExtensionExecutor
+
     @property
     @abstractmethod
     def priority(self) -> int:
@@ -77,6 +85,40 @@ def namespace(self) -> str:
     def validate(self, bot: Bot, event: Event) -> bool:
         return event.get_type() == "message"
 
+    @overload
+    async def inject(
+        self, dependent: Dependent[T], *, use_cache: bool = True, validate: bool | FieldInfo = False
+    ) -> T: ...
+
+    @overload
+    async def inject(self, dependent: tuple[str, type[T]]) -> T: ...
+
+    @overload
+    async def inject(self, dependent: Any) -> Any: ...
+
+    @final
+    async def inject(self, dependent: Any, use_cache: bool = True, validate: bool | FieldInfo = False) -> Any:
+        # assert isinstance(dependent, (Dependent, DependsInner)), "仅支持 Dependent 或 DependsInner 类型的依赖注入"
+        if isinstance(dependent, DependsInner):
+            if not dependent.dependency:
+                raise ValueError("DependsInner 未绑定任何依赖")
+            use_cache = dependent.use_cache
+            validate = dependent.validate
+            dependent = Dependent.parse(call=dependent.dependency, allow_types=self.executor.params)
+            param = DependParam(dependent=dependent, use_cache=use_cache, validate=validate)
+        elif isinstance(dependent, Dependent):
+            param = DependParam(dependent=dependent, use_cache=use_cache, validate=validate)
+        else:
+            for allow_type in self.executor.params:
+                if param := allow_type._check_param(
+                    inspect.Parameter(dependent[0], inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=dependent[1]),
+                    self.executor.params,
+                ):
+                    break
+            else:
+                raise ValueError(f"Unknown parameter {dependent[0]} with type {dependent[1]}")
+        return await self.executor._dependent_executor(param)
+
     async def output_converter(self, output_type: OutputType, content: str) -> UniMessage:
         """依据输出信息的类型，将字符串转换为消息对象以便发送。"""
         return UniMessage(content)
@@ -224,16 +266,39 @@ async def send_wrapper(self, bot: Bot, event: Event, send: TM) -> TM:
         return res
 
 
+class _DependentExecutor:
+    def __init__(
+        self,
+        bot: Bot,
+        event: Event,
+        state: T_State,
+        stack: AsyncExitStack | None = None,
+        dependency_cache: dict[_DependentCallable[Any], DependencyCache] | None = None,
+    ):
+        self.bot = bot
+        self.event = event
+        self.state = state
+        self.stack = stack
+        self.dependency_cache = dependency_cache or {}
+
+    async def __call__(self, param: Param):
+        return await param._solve(
+            stack=self.stack, dependency_cache=self.dependency_cache, bot=self.bot, event=self.event, state=self.state
+        )
+
+
 class ExtensionExecutor(SelectedExtensions):
     globals: ClassVar[list[type[Extension] | Extension]] = [DefaultExtension()]
     _rule: AlconnaRule
+    _dependent_executor: _DependentExecutor
 
     def __init__(
         self,
         rule: AlconnaRule,
         extensions: list[type[Extension] | Extension] | None = None,
         excludes: list[str | type[Extension]] | None = None,
     ):
+        self.params: tuple[type[Param], ...] = ()
         self.extensions: list[Extension] = []
         for ext in self.globals:
             if isinstance(ext, type):
@@ -284,6 +349,8 @@ def _callback(self, *append_global_ext: type[Extension] | Extension):
     def select(self, bot: Bot, event: Event) -> SelectedExtensions:
         context = [ext for ext in self.extensions if ext.validate(bot, event)]
         context.sort(key=lambda ext: ext.priority)
+        for ext in context:
+            ext.executor = self
         return SelectedExtensions(context)
 
     def before_catch(self, name: str, annotation: Any, default: Any) -> bool:

src/nonebot_plugin_alconna/matcher.py

@@ -1012,8 +1012,6 @@ def on_alconna(
         skip_for_unmatch,
         auto_send_output,
         comp_config,
-        extensions,
-        exclude_ext,
         use_origin,
         use_cmd_start,
         use_cmd_sep,
@@ -1022,13 +1020,14 @@ def on_alconna(
         rule,
         after_rule,
     )
-    executor = _rule.executor
-    params = (
+    _rule.executor = executor = ExtensionExecutor(_rule, extensions, exclude_ext)
+    executor.params = params = (
         ExtensionParam.new(executor),
         *Matcher.HANDLER_PARAM_TYPES[:-1],
         AlconnaParam,
         DefaultParam,
     )
+    executor.post_init(command)
     source = get_matcher_source(_depth + 1)
     NewMatcher = type(
         AlconnaMatcher.__name__,

src/nonebot_plugin_alconna/params.py

@@ -232,7 +232,7 @@ def _check_param(cls, param: inspect.Parameter, allow_types: tuple[type[Param],
             return cls(param.default, name=param.name, type=param.annotation, validate=True)
         return None
 
-    async def _solve(self, matcher: Matcher, event: Event, state: T_State, **kwargs: Any) -> Any:
+    async def _solve(self, event: Event, state: T_State, **kwargs: Any) -> Any:
         res = await self.executor.catch(event, state, self.extra["name"], self.extra["type"], self.default)
         if res is not PydanticUndefined:
             return res
@@ -275,7 +275,8 @@ def _check_param(cls, param: inspect.Parameter, allow_types: tuple[type[Param],
             return cls(..., type=Literal["context"])
         return cls(param.default, validate=True, name=param.name, type=param.annotation)
 
-    async def _solve(self, matcher: Matcher, event: Event, state: T_State, **kwargs: Any) -> Any:
+    async def _solve(self, event: Event, state: T_State, **kwargs: Any) -> Any:
+        matcher = kwargs.get("matcher", object())
         t = self.extra["type"]
         if ALCONNA_RESULT not in state:
             return self.default if self.default not in (..., Empty) else PydanticUndefined

src/nonebot_plugin_alconna/rule.py

@@ -18,7 +18,7 @@
 
 from .config import Config
 from .consts import ALCONNA_EXEC_RESULT, ALCONNA_EXTENSION, ALCONNA_RESULT, log
-from .extension import Extension, ExtensionExecutor, SelectedExtensions
+from .extension import ExtensionExecutor, SelectedExtensions, _DependentExecutor
 from .i18n import Lang
 from .model import CommandResult, CompConfig
 from .uniseg import UniMessage, UniMsg
@@ -58,6 +58,8 @@ class AlconnaRule:
         use_cmd_start: 是否使用 nb 全局配置里的命令前缀
     """
 
+    executor: ExtensionExecutor
+
     __slots__ = (
         "_comp_help",
         "_hide_tabs",
@@ -82,8 +84,6 @@ def __init__(
         skip_for_unmatch: bool = True,
         auto_send_output: Optional[bool] = None,
         comp_config: Optional[Union[CompConfig, bool]] = None,
-        extensions: Optional[list[Union[type[Extension], Extension]]] = None,
-        exclude_ext: Optional[list[Union[type[Extension], str]]] = None,
         use_origin: Optional[bool] = None,
         use_cmd_start: Optional[bool] = None,
         use_cmd_sep: Optional[bool] = None,
@@ -152,8 +152,6 @@ def _update(cmd_id: int):
             for alias in _aliases:
                 command.shortcut(alias, prefix=True, compact=None)
         self.skip = skip_for_unmatch
-        self.executor = ExtensionExecutor(self, extensions, exclude_ext)
-        self.executor.post_init(command)
         self._path = command.path
         self._namespace = command.namespace
         self._tasks: dict[str, asyncio.Task] = {}
@@ -297,11 +295,13 @@ async def __call__(
         stack: Optional[AsyncExitStack] = None,
         dependency_cache: Optional[dict[_DependentCallable[Any], DependencyCache]] = None,
     ) -> bool:
-        if not await self.before_rules(bot, event, state, stack, dependency_cache):
+
+        if self.before_rules.checkers and not await self.before_rules(bot, event, state, stack, dependency_cache):
             return False
         if event.get_type() == "meta_event":
             return False
         selected = self.executor.select(bot, event)
+        self.executor._dependent_executor = _DependentExecutor(bot, event, state, stack, dependency_cache)
         if not (msg := await selected.message_provider(event, state, bot, self.use_origin)):
             return False
         if not self.response_self and check_self_send(bot, event):
@@ -356,13 +356,13 @@ async def __call__(
             return False
         if not await selected.permission_check(bot, event, cmd):
             return False
-        await selected.parse_wrapper(bot, state, event, arp)
         state[ALCONNA_RESULT] = CommandResult(result=arp, output=may_help_text)
         state[ALCONNA_EXEC_RESULT] = cmd.exec_result
         state[ALCONNA_EXTENSION] = selected
-        if not await self.after_rules(bot, event, state, stack, dependency_cache):
-            return False
-        return True
+        await selected.parse_wrapper(bot, state, event, arp)
+        return not (
+            self.after_rules.checkers and not await self.after_rules(bot, event, state, stack, dependency_cache)
+        )
 
     async def send(self, text: str, bot: Bot, event: Event, arp: Arparma) -> Any:
         _t = str(arp.error_info) if isinstance(arp.error_info, SpecialOptionTriggered) else "error"

src/nonebot_plugin_alconna/uniseg/message.py

@@ -6,7 +6,7 @@
 from json import dumps, loads
 from pathlib import Path
 from types import FunctionType
-from typing import TYPE_CHECKING, Any, Callable, Literal, NoReturn, TypeVar, Union, Protocol
+from typing import TYPE_CHECKING, Any, Callable, Literal, NoReturn, Protocol, TypeVar, Union
 from typing_extensions import Self, SupportsIndex, deprecated
 
 from nonebot.exception import FinishedException

tests/test_extension.py

@@ -1,9 +1,10 @@
 from typing import Any
 
-from arclet.alconna import Alconna, Args
+from arclet.alconna import Alconna, Args, Arparma
 from nonebot import get_adapter
-from nonebot.adapters.onebot.v11 import Adapter, Bot, Message
 from nonebot.internal.adapter import Event
+from nonebot.params import Depends
+from nonebot.typing import T_State
 from nonebug import App
 import pytest
 
@@ -12,9 +13,18 @@
 
 @pytest.mark.asyncio()
 async def test_extension(app: App):
-    from nonebot.adapters.onebot.v11 import MessageEvent
+    from nonebot.adapters.onebot.v11 import Adapter, Bot, Message, MessageEvent
 
     from nonebot_plugin_alconna import Extension, Interface, UniMessage, on_alconna
+    from nonebot_plugin_alconna.params import Match
+
+    async def perm_check(bot, event):
+        if event.get_user_id() != "123":
+            await bot.send(event, "权限不足！")
+            return False
+        return True
+
+    check_dep = Depends(perm_check, use_cache=True)
 
     class DemoExtension(Extension):
         @property
@@ -26,10 +36,7 @@ def id(self) -> str:
             return "demo"
 
         async def permission_check(self, bot, event, command):
-            if event.get_user_id() != "123":
-                await bot.send(event, "权限不足！")
-                return False
-            return True
+            return await self.inject(check_dep)
 
         def before_catch(self, name: str, annotation: Any, default: Any) -> bool:
             return annotation is str
@@ -47,6 +54,12 @@ async def catch(self, interface: Interface[MessageEvent]):
                 }.get(interface.name, interface.name)
             return None
 
+        async def parse_wrapper(self, bot: Bot, state: T_State, event: Event, res: Arparma) -> None:
+            a = await self.inject(("a", Match[float]))
+            b = await self.inject(("b", Match[float]))
+            assert a.result == 1.3
+            assert b.result == 2.4
+
     add = on_alconna(Alconna("add", Args["a", float]["b", float]), extensions=[DemoExtension], comp_config={})
 
     @add.handle()