Moved cooldowns to PosterUser class

classes/user.py

@@ -1,5 +1,7 @@
+from modules.app import app
 from datetime import datetime
-from modules.database import col_banned
+from modules.database import col_banned, col_users
+from modules.utils import configGet
 
 class PosterUser():
 
@@ -15,3 +17,16 @@ class PosterUser():
 
     def unblock(self) -> None:
         col_banned.find_one_and_delete({"user": self.id})
+
+    def is_limited(self) -> bool:
+        if self.id in app.admins:
+            return False
+        else:
+            db_record = col_users.find_one({"user": self.id})
+            if db_record is None:
+                return False
+            return True if (datetime.now() - db_record["cooldown"]).total_seconds() < configGet("timeout", "submission") else False
+        
+    def limit(self) -> None:
+        if col_users.find_one_and_update({"user": self.id}, {"$set": {"cooldown": datetime.now()}}) is None:
+            col_users.insert_one({"user": self.id, "cooldown": datetime.now()})

modules/submissions.py

@@ -1,18 +0,0 @@
-from modules.app import app
-from datetime import datetime
-from modules.utils import configGet
-from modules.database import col_users
-from pyrogram.types.user_and_chats import User
-
-def subLimit(user: User) -> None:
-    if col_users.find_one_and_update({"user": user.id}, {"$set": {"cooldown": datetime.now()}}) is None:
-        col_users.insert_one({"user": user.id, "cooldown": datetime.now()})
-
-def subLimited(user: User) -> bool:
-    if user.id in app.admins:
-        return False
-    else:
-        db_record = col_users.find_one({"user": user.id})
-        if db_record is None:
-            return False
-        return True if (datetime.now() - db_record["cooldown"]).total_seconds() < configGet("timeout", "submission") else False

plugins/handlers/submission.py

@@ -1,4 +1,4 @@
-from datetime import datetime, timezone
+from datetime import datetime
 from os import makedirs, path, sep
 from traceback import format_exc
 from uuid import uuid4
@@ -8,11 +8,11 @@ from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup, Message
 from pyrogram.enums.chat_action import ChatAction
 from classes.exceptions import SubmissionDuplicatesError
 from classes.poster_client import PosterClient
+from classes.user import PosterUser
 
 from modules.app import app
 from modules.database import col_banned, col_submitted
 from modules.logger import logWrite
-from modules.submissions import subLimit, subLimited
 from modules.utils import configGet, locale
 from classes.enums.submission_types import SubmissionType
 
@@ -31,7 +31,7 @@ async def get_submission(app: PosterClient, msg: Message):
         save_tmp = True
         contents = None
 
-        if subLimited(msg.from_user):
+        if PosterUser(msg.from_user.id).is_limited():
             await msg.reply_text(locale("sub_cooldown", "message", locale=user_locale).format(str(configGet("timeout", "submission"))))
             return
 
@@ -180,7 +180,7 @@ async def get_submission(app: PosterClient, msg: Message):
                 # ]
             ]
 
-            subLimit(msg.from_user)
+            PosterUser(msg.from_user.id).limit()
 
         if msg.from_user.id != app.owner:
             await msg.reply_text(locale("sub_sent", "message", locale=user_locale), quote=True)