优化manager代码

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
+106 -61
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 FreqLimiter
from services.log import logger
from pathlib import Path
from ruamel.yaml import YAML
from ruamel import yaml
from .models import PluginCd
yaml = YAML(typ="safe")
_yaml = yaml.YAML(typ="safe")
class Plugins2cdManager(StaticData):
@@ -14,16 +15,9 @@ class Plugins2cdManager(StaticData):
"""
def __init__(self, file: Path):
self.file = file
super().__init__(None)
super().__init__(file, False)
self._freq_limiter: Dict[str, FreqLimiter] = {}
if file.exists():
with open(file, "r", encoding="utf8") as f:
self._data = yaml.load(f)
if "PluginCdLimit" in self._data.keys():
self._data = (
self._data["PluginCdLimit"] if self._data["PluginCdLimit"] else {}
)
self.__load_file()
def add_cd_limit(
self,
@@ -31,42 +25,35 @@ class Plugins2cdManager(StaticData):
*,
cd: Optional[int] = 5,
status: Optional[bool] = True,
check_type: Optional[str] = "all",
limit_type: Optional[str] = "user",
check_type: Literal["private", "group", "all"] = "all",
limit_type: Literal["user", "group"] = "user",
rst: Optional[str] = None,
**kwargs # 用于接收额外实参
):
"""
添加插件调用 cd 限制
:param plugin: 插件模块名称
:param cd: cd 时长
:param status: 默认开关状态
:param check_type: 检查类型 'private'/'group'/'all',限制私聊/群聊/全部
:param limit_type: 限制类型 监听对象,以user_id或group_id作为键来限制,'user':用户id,'group':群id
:param rst: 回复的话,为空则不回复
说明:
添加插件调用 cd 限制
参数:
:param plugin: 插件模块名称
:param cd: cd 时长
:param status: 默认开关状态
:param check_type: 检查类型 'private'/'group'/'all',限制私聊/群聊/全部
:param limit_type: 限制类型 监听对象,以user_id或group_id作为键来限制,'user':用户id,'group':群id
:param rst: 回复的话,为空则不回复
"""
cd = cd or 5
status = status or True
check_type = check_type or "all"
limit_type = limit_type or "user"
if check_type not in ["all", "group", "private"]:
raise ValueError(
f"{plugin} 添加cd限制错误,‘check_type‘ 必须为 'private'/'group'/'all'"
)
if limit_type not in ["user", "group"]:
raise ValueError(f"{plugin} 添加cd限制错误,‘limit_type‘ 必须为 'user'/'group'")
self._data[plugin] = {
"cd": cd,
"status": status,
"check_type": check_type,
"limit_type": limit_type,
"rst": rst,
}
self._data[plugin] = PluginCd(cd=cd, status=status, check_type=check_type, limit_type=limit_type, rst=rst)
def get_plugin_cd_data(self, plugin: str) -> Optional[dict]:
def get_plugin_cd_data(self, plugin: str) -> Optional[PluginCd]:
"""
获取插件cd数据
:param plugin: 模块名
说明:
获取插件cd数据
参数:
:param plugin: 模块名
"""
if self.check_plugin_cd_status(plugin):
return self._data[plugin]
@@ -74,20 +61,24 @@ class Plugins2cdManager(StaticData):
def check_plugin_cd_status(self, plugin: str) -> bool:
"""
检测插件是否有 cd
:param plugin: 模块名
说明:
检测插件是否有 cd
参数:
:param plugin: 模块名
"""
return (
plugin in self._data.keys()
and self._data[plugin]["cd"] > 0
and self._data[plugin]["status"]
and self._data[plugin].cd > 0
and self._data[plugin].status
)
def check(self, plugin: str, id_: int) -> bool:
"""
检查 cd
:param plugin: 模块名
:param id_: 限制 id
说明:
检查 cd
参数:
:param plugin: 模块名
:param id_: 限制 id
"""
if self._freq_limiter.get(plugin):
return self._freq_limiter[plugin].check(id_)
@@ -95,42 +86,96 @@ class Plugins2cdManager(StaticData):
def start_cd(self, plugin: str, id_: int, cd: int = 0):
"""
开始cd
:param plugin: 模块名
:param id_: cd 限制类型
:param cd: cd 时长
:return:
说明:
开始cd
参数:
:param plugin: 模块名
:param id_: cd 限制类型
:param cd: cd 时长
"""
if self._freq_limiter.get(plugin):
self._freq_limiter[plugin].start_cd(id_, cd)
def get_plugin_data(self, plugin: str) -> dict:
def get_plugin_data(self, plugin: str) -> Optional[PluginCd]:
"""
获取单个模块限制数据
:param plugin: 模块名
说明:
获取单个模块限制数据
参数:
:param plugin: 模块名
"""
if self._data.get(plugin) is not None:
if self._data.get(plugin):
return self._data.get(plugin)
return {}
def reload_cd_limit(self):
"""
加载 cd 限制器
:return:
说明:
加载 cd 限制器
"""
for plugin in self._data:
if self.check_plugin_cd_status(plugin):
self._freq_limiter[plugin] = FreqLimiter(
self.get_plugin_cd_data(plugin)["cd"]
self.get_plugin_cd_data(plugin).cd
)
logger.info(f"已成功加载 {len(self._freq_limiter)} 个Cd限制.")
def reload(self):
"""
重载本地数据
说明:
重载本地数据
"""
self.__load_file()
self.reload_cd_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(
{"PluginCdLimit": self.dict()},
f,
indent=2,
Dumper=yaml.RoundTripDumper,
allow_unicode=True,
)
_data = yaml.round_trip_load(open(path, encoding="utf8"))
_data["PluginCdLimit"].yaml_set_start_comment(
"""# 需要cd的功能
# 自定义的功能需要cd也可以在此配置
# key:模块名称
# cd:cd 时长(秒)
# status:此限制的开关状态
# check_type:'private'/'group'/'all',限制私聊/群聊/全部
# limit_type:监听对象,以user_id或group_id作为键来限制,'user':用户id,'group':群id
# 示例:'user':用户N秒内触发1次,'group':群N秒内触发1次
# 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, PluginCd] = {}
if self.file.exists():
with open(self.file, "r", encoding="utf8") as f:
self._data: dict = yaml.load(f)
self._data = self._data["PluginCdLimit"]
self.reload_cd_limit()
temp = _yaml.load(f)
if "PluginCdLimit" in temp.keys():
for k, v in temp["PluginCdLimit"].items():
self._data[k] = PluginCd.parse_obj(v)