refactor: optimize MuteManager data source

zhenxun/plugins/mute/_data_source.py

@@ -1,6 +1,6 @@
 import time
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 import ujson as json
 
 from zhenxun.configs.config import Config
@@ -16,7 +16,7 @@ class GroupData(BaseModel):
     """检测时长"""
     duration: int
     """禁言时长"""
-    message_data: dict = {}
+    message_data: dict = Field(default_factory=dict)
     """消息存储"""
 
 
@@ -26,12 +26,13 @@ class MuteManager:
     def __init__(self) -> None:
         self._group_data: dict[str, GroupData] = {}
         if self.file.exists():
-            _data = json.load(open(self.file))
-            for gid in _data:
+            with open(self.file, encoding="utf-8") as f:
+                _data = json.load(f)
+            for gid, gdata in _data.items():
                 self._group_data[gid] = GroupData(
-                    count=_data[gid]["count"],
-                    time=_data[gid]["time"],
-                    duration=_data[gid]["duration"],
+                    count=gdata["count"],
+                    time=gdata["time"],
+                    duration=gdata["duration"],
                 )
 
     def get_group_data(self, group_id: str) -> GroupData:
@@ -64,14 +65,15 @@ class MuteManager:
 
     def save_data(self):
         """保存数据"""
-        data = {}
-        for gid in self._group_data:
-            data[gid] = {
-                "count": self._group_data[gid].count,
-                "time": self._group_data[gid].time,
-                "duration": self._group_data[gid].duration,
+        data = {
+            gid: {
+                "count": gdata.count,
+                "time": gdata.time,
+                "duration": gdata.duration,
             }
-        with open(self.file, "w") as f:
+            for gid, gdata in self._group_data.items()
+        }
+        with open(self.file, "w", encoding="utf-8") as f:
             json.dump(data, f, indent=4, ensure_ascii=False)
 
     def add_message(self, user_id: str, group_id: str, message: str) -> int:
@@ -85,37 +87,42 @@ class MuteManager:
         返回:
             int: 禁言时长
         """
-        if group_id not in self._group_data:
-            self._group_data[group_id] = GroupData(
-                count=base_config.get("MUTE_DEFAULT_COUNT"),
-                time=base_config.get("MUTE_DEFAULT_TIME"),
-                duration=base_config.get("MUTE_DEFAULT_DURATION"),
-            )
-        group_data = self._group_data[group_id]
+        group_data = self.get_group_data(group_id)
         if group_data.duration == 0:
             return 0
+
         message_data = group_data.message_data
-        if not message_data.get(user_id):
+        user_data = message_data.get(user_id)
+        now = time.time()
+
+        if not user_data:
             message_data[user_id] = {
-                "time": time.time(),
+                "time": now,
                 "count": 1,
                 "message": message,
             }
+            return 0
+
+        # 超过检测时间窗口，重置计数
+        if now - user_data["time"] > group_data.time:
+            user_data["time"] = now
+            user_data["count"] = 1
+            user_data["message"] = message
+            return 0
+
+        # 消息内容相似（包含之前的消息），累加计数
+        if user_data["message"] in message:
+            user_data["count"] += 1
         else:
-            if message.find(message_data[user_id]["message"]) != -1:
-                message_data[user_id]["count"] += 1
-            else:
-                message_data[user_id]["time"] = time.time()
-                message_data[user_id]["count"] = 1
-            message_data[user_id]["message"] = message
-            if time.time() - message_data[user_id]["time"] > group_data.time:
-                message_data[user_id]["time"] = time.time()
-                message_data[user_id]["count"] = 1
-            if (
-                message_data[user_id]["count"] > group_data.count
-                and time.time() - message_data[user_id]["time"] < group_data.time
-            ):
-                return group_data.duration
+            user_data["time"] = now
+            user_data["count"] = 1
+
+        user_data["message"] = message
+
+        # 检测是否触发刷屏
+        if user_data["count"] > group_data.count:
+            return group_data.duration
+
         return 0