Merge pull request #57 from iAliF/main

Add 'admin_only' decorator

tgbot/handlers/admin/handlers.py

@@ -6,26 +6,19 @@
 
 from tgbot.handlers.admin import static_text
 from tgbot.handlers.admin.utils import _get_csv_from_qs_values
-from tgbot.handlers.utils.info import send_typing_action
+from tgbot.handlers.utils.decorators import admin_only, send_typing_action
 from users.models import User
 
 
+@admin_only
 def admin(update: Update, context: CallbackContext) -> None:
     """ Show help info about all secret admins commands """
-    u = User.get_user(update, context)
-    if not u.is_admin:
-        update.message.reply_text(static_text.only_for_admins)
-        return
     update.message.reply_text(static_text.secret_admin_commands)
 
 
+@admin_only
 def stats(update: Update, context: CallbackContext) -> None:
     """ Show help info about all secret admins commands """
-    u = User.get_user(update, context)
-    if not u.is_admin:
-        update.message.reply_text(static_text.only_for_admins)
-        return
-
     text = static_text.users_amount_stat.format(
         user_count=User.objects.count(),  # count may be ineffective if there are a lot of users.
         active_24=User.objects.filter(updated_at__gte=now() - timedelta(hours=24)).count()
@@ -38,14 +31,10 @@ def stats(update: Update, context: CallbackContext) -> None:
     )
 
 
+@admin_only
 @send_typing_action
 def export_users(update: Update, context: CallbackContext) -> None:
-    u = User.get_user(update, context)
-    if not u.is_admin:
-        update.message.reply_text(static_text.only_for_admins)
-        return
-
     # in values argument you can specify which fields should be returned in output csv
     users = User.objects.all().values()
     csv_users = _get_csv_from_qs_values(users)
-    context.bot.send_document(chat_id=u.user_id, document=csv_users)
+    update.message.reply_document(csv_users)

tgbot/handlers/admin/static_text.py

@@ -1,5 +1,4 @@
 command_start = '/stats'
-only_for_admins = 'Sorry, this function is available only for admins. Set "admin" flag in django admin panel.'
 
 secret_admin_commands = f"⚠️ Secret Admin commands\n" \
                         f"{command_start} - bot stats"

tgbot/handlers/utils/decorators.py

@@ -0,0 +1,36 @@
+from functools import wraps
+from typing import Callable
+
+from telegram import Update, ChatAction
+from telegram.ext import CallbackContext
+
+from users.models import User
+
+
+def admin_only(func: Callable):
+    """
+    Admin only decorator
+    Used for handlers that only admins have access to
+    """
+
+    @wraps(func)
+    def wrapper(update: Update, context: CallbackContext, *args, **kwargs):
+        user = User.get_user(update, context)
+
+        if not user.is_admin:
+            return
+
+        return func(update, context, *args, **kwargs)
+
+    return wrapper
+
+
+def send_typing_action(func: Callable):
+    """Sends typing action while processing func command."""
+
+    @wraps(func)
+    def command_func(update: Update, context: CallbackContext, *args, **kwargs):
+        update.effective_chat.send_chat_action(ChatAction.TYPING)
+        return func(update, context, *args, **kwargs)
+
+    return command_func

tgbot/handlers/utils/info.py

@@ -1,22 +1,7 @@
-from functools import wraps
-from typing import Dict, Callable
+from typing import Dict
 
-import telegram
 from telegram import Update
 
-from tgbot.main import bot
-
-
-def send_typing_action(func: Callable):
-    """Sends typing action while processing func command."""
-    @wraps(func)
-    def command_func(update, context, *args, **kwargs):
-        bot.send_chat_action(
-            chat_id=update.effective_message.chat_id, action=telegram.ChatAction.TYPING)
-        return func(update, context,  *args, **kwargs)
-
-    return command_func
-
 
 def extract_user_data_from_update(update: Update) -> Dict:
     """ python-telegram-bot's Update instance --> User info """