更新至nonebot-rc1

basic_plugins/shop/use/data_source.py

@@ -3,7 +3,7 @@ from services.log import logger
 from nonebot.adapters.onebot.v11 import Bot
 from pydantic import create_model
 from utils.models import ShopParam
-from typing import Optional, Union
+from typing import Optional, Union, Callable, List, Tuple, Dict, Any
 from types import MappingProxyType
 import inspect
 import asyncio
@@ -13,6 +13,30 @@ class GoodsUseFuncManager:
     def __init__(self):
         self._data = {}
 
+    def register_use_before_handle(self, goods_name: str, fun_list: List[Callable]):
+        """
+        说明:
+            注册商品使用前函数
+        参数:
+            :param goods_name: 商品名称
+            :param fun_list: 函数列表
+        """
+        if fun_list:
+            self._data[goods_name]["before_handle"] = fun_list
+            logger.info(f"register_use_before_handle 成功注册商品：{goods_name} 的{len(fun_list)}个使用前函数")
+
+    def register_use_after_handle(self, goods_name: str, fun_list: List[Callable]):
+        """
+        说明:
+            注册商品使用后函数
+        参数:
+            :param goods_name: 商品名称
+            :param fun_list: 函数列表
+        """
+        if fun_list:
+            self._data[goods_name]["after_handle"] = fun_list
+            logger.info(f"register_use_after_handle 成功注册商品：{goods_name} 的{len(fun_list)}个使用后函数")
+
     def register_use(self, goods_name: str, **kwargs):
         """
         注册商品使用方法
@@ -26,7 +50,7 @@ class GoodsUseFuncManager:
         判断商品使用方法是否被注册
         :param goods_name: 商品名称
         """
-        return bool(self ._data.get(goods_name))
+        return bool(self._data.get(goods_name))
 
     def get_max_num_limit(self, goods_name: str) -> int:
         """
@@ -37,6 +61,21 @@ class GoodsUseFuncManager:
             return self._data[goods_name]["kwargs"]["max_num_limit"]
         return 1
 
+    def _parse_args(self, args_: MappingProxyType, param: ShopParam, **kwargs):
+        param_list_ = []
+        _bot = param.bot
+        param.bot = None
+        param_json = param.dict()
+        param_json["bot"] = _bot
+        for par in args_.keys():
+            if par in ["shop_param"]:
+                param_list_.append(param)
+            elif par not in ["args", "kwargs"]:
+                param_list_.append(param_json.get(par))
+                if kwargs.get(par) is not None:
+                    del kwargs[par]
+        return param_list_
+
     async def use(
         self, param: ShopParam, **kwargs
     ) -> Optional[Union[str, MessageSegment]]:
@@ -45,31 +84,18 @@ class GoodsUseFuncManager:
         :param param: BaseModel
         :param kwargs: kwargs
         """
-        def parse_args(args_: MappingProxyType):
-            param_list_ = []
-            _bot = param.bot
-            param.bot = None
-            param_json = param.dict()
-            param_json["bot"] = _bot
-            for par in args_.keys():
-                if par in ["shop_param"]:
-                    param_list_.append(param)
-                elif par not in ["args", "kwargs"]:
-                    param_list_.append(param_json.get(par))
-                    if kwargs.get(par) is not None:
-                        del kwargs[par]
-            return param_list_
         goods_name = param.goods_name
         if self.exists(goods_name):
+            # 使用方法
             args = inspect.signature(self._data[goods_name]["func"]).parameters
             if args and list(args.keys())[0] != "kwargs":
                 if asyncio.iscoroutinefunction(self._data[goods_name]["func"]):
                     return await self._data[goods_name]["func"](
-                        *parse_args(args)
+                        *self._parse_args(args, param, **kwargs)
                     )
                 else:
                     return self._data[goods_name]["func"](
-                        *parse_args(args)
+                        *self._parse_args(args, param, **kwargs)
                     )
             else:
                 if asyncio.iscoroutinefunction(self._data[goods_name]["func"]):
@@ -81,6 +107,21 @@ class GoodsUseFuncManager:
                         **kwargs,
                     )
 
+    async def run_handle(self, goods_name: str, type_: str, param: ShopParam, **kwargs):
+        if self._data[goods_name].get(type_):
+            for func in self._data[goods_name].get(type_):
+                args = inspect.signature(func).parameters
+                if args and list(args.keys())[0] != "kwargs":
+                    if asyncio.iscoroutinefunction(func):
+                        await func(*self._parse_args(args, param, **kwargs))
+                    else:
+                        func(*self._parse_args(args, param, **kwargs))
+                else:
+                    if asyncio.iscoroutinefunction(func):
+                        await func(**kwargs)
+                    else:
+                        func(**kwargs)
+
     def check_send_success_message(self, goods_name: str) -> bool:
         """
         检查是否发送使用成功信息
@@ -115,6 +156,34 @@ class GoodsUseFuncManager:
 func_manager = GoodsUseFuncManager()
 
 
+def build_params(
+    bot: Bot, event: GroupMessageEvent, goods_name: str, num: int
+) -> Tuple[ShopParam, Dict[str, Any]]:
+    """
+    说明:
+        构造参数
+    参数:
+        :param bot: bot
+        :param event: event
+        :param goods_name: 商品名称
+        :param num: 数量
+    :return:
+    """
+    _kwargs = func_manager.get_kwargs(goods_name)
+    return (
+        func_manager.init_model(goods_name, bot, event, num),
+        {
+            **_kwargs,
+            "_bot": bot,
+            "event": event,
+            "group_id": event.group_id,
+            "user_id": event.user_id,
+            "num": num,
+            "goods_name": goods_name,
+        },
+    )
+
+
 async def effect(
     bot: Bot, event: GroupMessageEvent, goods_name: str, num: int
 ) -> Optional[Union[str, MessageSegment]]:
@@ -130,24 +199,14 @@ async def effect(
     # try:
     if func_manager.exists(goods_name):
         _kwargs = func_manager.get_kwargs(goods_name)
-        return await func_manager.use(
-            func_manager.init_model(goods_name, bot, event, num),
-            **{
-                **_kwargs,
-                "_bot": bot,
-                "event": event,
-                "group_id": event.group_id,
-                "user_id": event.user_id,
-                "num": num,
-                "goods_name": goods_name,
-            },
-        )
+        model, kwargs = build_params(bot, event, goods_name, num)
+        return await func_manager.use(model, **kwargs)
     # except Exception as e:
     #     logger.error(f"use 商品生效函数effect 发生错误 {type(e)}：{e}")
     return None
 
 
-def register_use(goods_name: str, func, **kwargs):
+def register_use(goods_name: str, func: Callable, **kwargs):
     """
     注册商品使用方法
     :param goods_name: 商品名称