🎨 新增冷却组件

1 files changed, 32 insertions, 0 deletions

diff --git a/ATRI/utils/__init__.py b/ATRI/utils/__init__.py
index aeeb9db..a1a7f3b 100644
--- a/ATRI/utils/__init__.py
+++ b/ATRI/utils/__init__.py
@@ -5,11 +5,13 @@ import pytz
 import yaml
 import time
 import string
+import asyncio
 import aiofiles
 from pathlib import Path
 from random import sample
 from datetime import datetime
 from PIL import Image, ImageFile
+from collections import defaultdict
 from aiofiles.threadpool.text import AsyncTextIOWrapper
 
 
@@ -220,3 +222,33 @@ class Translate:
                 output_str_list.append(self.text[i])
 
         return "".join(output_str_list)
+
+
+class Limiter:
+    def __init__(self, max_count: int, down_time: float):
+        """冷却设置
+
+        Args:
+            max_count (int): 最大次数
+            down_time (float): 到达次数后的冷却时间
+        """
+        self.max_count = max_count
+        self.down_time = down_time
+        self.count = defaultdict(int)
+
+    def check(self, key: str) -> bool:
+        if self.count[key] >= self.max_count:
+            loop = asyncio.get_running_loop()
+            loop.call_later(self.down_time, self.reset)
+            return False
+            
+        return True
+
+    def increase(self, key: str, times: int = 1) -> None:
+        self.count[key] += times
+
+    def reset(self, key: str) -> None:
+        self.count[key] = 0
+
+    def get_times(self, key: str) -> int:
+        return self.count[key]