From fef46d6fb71df12e35068cad5d83da5e070629aa Mon Sep 17 00:00:00 2001
From: Kyomotoi <0w0@imki.moe>
Date: Fri, 2 Dec 2022 08:54:23 +0800
Subject: =?UTF-8?q?=F0=9F=8E=A8=20=E6=96=B0=E5=A2=9E=E5=86=B7=E5=8D=B4?=
 =?UTF-8?q?=E7=BB=84=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ATRI/utils/__init__.py | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

(limited to 'ATRI')

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]
-- 
cgit v1.2.3