Diuse/zhenxun_bot
1
0
Fork 0
mirror of https://github.com/zhenxun-org/zhenxun_bot.git synced 2026-09-29 16:50:02 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

优化manager代码

Browse Source
This commit is contained in:
HibiKier
2022-11-21 20:43:41 +08:00
parent 49e38862af
commit cd66b7f5b8
41 changed files with 1974 additions and 1760 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/manager/plugins2count_manager.py
Executable → Regular
+101 -55
View File
@@ -1,11 +1,12 @@
from typing import Optional, Dict
from .data_class import StaticData
from typing import Optional, Dict, Literal, Union
from utils.manager.data_class import StaticData
from utils.utils import DailyNumberLimiter
from services.log import logger
from pathlib import Path
from ruamel.yaml import YAML
from ruamel import yaml
from .models import PluginCount
yaml = YAML(typ="safe")
_yaml = yaml.YAML(typ="safe")
class Plugins2countManager(StaticData):
@@ -14,16 +15,9 @@ class Plugins2countManager(StaticData):
"""
def __init__(self, file: Path):
self.file = file
super().__init__(None)
super().__init__(file, False)
self._daily_limiter: Dict[str, DailyNumberLimiter] = {}
if file.exists():
with open(file, "r", encoding="utf8") as f:
self._data = yaml.load(f)
if "PluginCountLimit" in self._data.keys():
self._data = (
self._data["PluginCountLimit"] if self._data["PluginCountLimit"] else {}
)
self.__load_file()
def add_count_limit(
self,
@@ -31,64 +25,64 @@ class Plugins2countManager(StaticData):
*,
max_count: int = 5,
status: Optional[bool] = True,
limit_type: Optional[str] = "user",
limit_type: Literal["user", "group"] = "user",
rst: Optional[str] = None,
**kwargs # 用于接收额外实参
):
"""
添加插件调用 次数 限制
:param plugin: 插件模块名称
:param max_count: 最大次数限制
:param status: 默认开关状态
:param limit_type: 限制类型 监听对象，以user_id或group_id作为键来限制，'user'：用户id，'group'：群id
:param rst: 回复的话，为空则不回复
说明:
添加插件调用 次数 限制
参数:
:param plugin: 插件模块名称
:param max_count: 最大次数限制
:param status: 默认开关状态
:param limit_type: 限制类型 监听对象，以user_id或group_id作为键来限制，'user'：用户id，'group'：群id
:param rst: 回复的话，为空则不回复
"""
max_count = max_count or 5
status = status or True
limit_type = limit_type or "user"
if limit_type not in ["user", "group"]:
raise ValueError(f"{plugin} 添加count限制错误，‘limit_type‘ 必须为 'user'/'group'")
self._data[plugin] = {
"max_count": max_count,
"status": status,
"limit_type": limit_type,
"rst": rst,
}
self._data[plugin] = PluginCount(max_count=max_count, status=status, limit_type=limit_type, rst=rst)
def get_plugin_count_data(self, plugin: str) -> Optional[dict]:
def get_plugin_count_data(self, plugin: str) -> Optional[PluginCount]:
"""
获取插件次数数据
:param plugin: 模块名
说明:
获取插件次数数据
参数:
:param plugin: 模块名
"""
if self.check_plugin_count_status(plugin):
return self._data[plugin]
return None
def get_plugin_data(self, plugin: str) -> Optional[dict]:
def get_plugin_data(self, plugin: str) -> Optional[PluginCount]:
"""
获取单个模块限制数据
:param plugin: 模块名
说明:
获取单个模块限制数据
参数:
:param plugin: 模块名
"""
if self._data.get(plugin) is not None:
return self._data.get(plugin)
return None
def check_plugin_count_status(self, plugin: str) -> bool:
"""
检测插件是否有 次数 限制
:param plugin: 模块名
说明:
检测插件是否有 次数 限制
参数:
:param plugin: 模块名
"""
return (
plugin in self._data.keys()
and self._data[plugin]["status"]
and self._data[plugin]["max_count"] > 0
and self._data[plugin].status
and self._data[plugin].max_count > 0
)
def check(self, plugin: str, id_: int) -> bool:
"""
检查 count
:param plugin: 模块名
:param id_: 限制 id
说明:
检查 count
参数:
:param plugin: 模块名
:param id_: 限制 id
"""
if self._daily_limiter.get(plugin):
return self._daily_limiter[plugin].check(id_)
@@ -96,24 +90,25 @@ class Plugins2countManager(StaticData):
def increase(self, plugin: str, id_: int, num: int = 1):
"""
增加次数
:param plugin: 模块名
:param id_: cd 限制类型
:param num: 增加次数
:return:
说明:
增加次数
参数:
:param plugin: 模块名
:param id_: cd 限制类型
:param num: 增加次数
"""
if self._daily_limiter.get(plugin):
self._daily_limiter[plugin].increase(id_, num)
def reload_count_limit(self):
"""
加载 cd 限制器
:return:
说明:
加载 cd 限制器
"""
for plugin in self._data:
if self.check_plugin_count_status(plugin):
self._daily_limiter[plugin] = DailyNumberLimiter(
self.get_plugin_count_data(plugin)["max_count"]
self.get_plugin_count_data(plugin).max_count
)
logger.info(f"已成功加载 {len(self._daily_limiter)} 个Count限制.")
@@ -121,8 +116,59 @@ class Plugins2countManager(StaticData):
"""
重载本地数据
"""
self.__load_file()
self.reload_count_limit()
def save(self, path: Union[str, Path] = None):
"""
说明:
保存文件
参数:
:param path: 文件路径
"""
path = path or self.file
if isinstance(path, str):
path = Path(path)
if path:
with open(path, "w", encoding="utf8") as f:
yaml.dump(
{"PluginCountLimit": self.dict()},
f,
indent=2,
Dumper=yaml.RoundTripDumper,
allow_unicode=True,
)
_data = yaml.round_trip_load(open(path, encoding="utf8"))
_data["PluginCountLimit"].yaml_set_start_comment(
"""# 命令每日次数限制
# 即 用户/群聊 每日可调用命令的次数 [数据内存存储，重启将会重置]
# 每日调用直到 00:00 刷新
# key：模块名称
# max_count: 每日调用上限
# status：此限制的开关状态
# limit_type：监听对象，以user_id或group_id作为键来限制，'user'：用户id，'group'：群id
# 示例：'user'：用户上限，'group'：群聊上限
# rst：回复的话，可以添加[at]，[uname]，[nickname]来对应艾特，用户群名称，昵称系统昵称
# rst 为 "" 或 None 时则不回复
# rst示例："[uname]你冲的太快了，[nickname]先生，请稍后再冲[at]"
# rst回复："老色批你冲的太快了，欧尼酱先生，请稍后再冲@老色批"
# 用户昵称↑ 昵称系统的昵称↑ 艾特用户↑""",
indent=2,
)
with open(path, "w", encoding="utf8") as wf:
yaml.round_trip_dump(
_data, wf, Dumper=yaml.RoundTripDumper, allow_unicode=True
)
def __load_file(self):
"""
说明:
读取配置文件
"""
self._data: Dict[str, PluginCount] = {}
if self.file.exists():
with open(self.file, "r", encoding="utf8") as f:
self._data: dict = yaml.load(f)
self._data = self._data["PluginCountLimit"]
self.reload_count_limit()
temp = _yaml.load(f)
if "PluginCountLimit" in temp.keys():
for k, v in temp["PluginCountLimit"].items():
self._data[k] = PluginCount.parse_obj(v)