From 3fb6980095d14728f3697dc99be734f6ab08e9cd Mon Sep 17 00:00:00 2001 From: Mohammed Rabil Date: Thu, 28 Dec 2023 16:42:44 +0530 Subject: [PATCH] Refactor filter names and imports in bot.py --- gemini_pro_bot/bot.py | 17 ++++++----------- gemini_pro_bot/filters.py | 20 ++++++++++++++++---- 2 files changed, 22 insertions(+), 15 deletions(-) diff --git a/gemini_pro_bot/bot.py b/gemini_pro_bot/bot.py index e8eded0..3bf816d 100644 --- a/gemini_pro_bot/bot.py +++ b/gemini_pro_bot/bot.py @@ -3,10 +3,9 @@ from telegram.ext import ( CommandHandler, MessageHandler, - filters, Application, ) -from gemini_pro_bot.filters import AuthorizedUserFilter +from gemini_pro_bot.filters import AuthFilter, MessageFilter, PhotoFilter from dotenv import load_dotenv from gemini_pro_bot.handlers import ( start, @@ -25,19 +24,15 @@ def start_bot() -> None: application = Application.builder().token(os.getenv("BOT_TOKEN")).build() # on different commands - answer in Telegram - application.add_handler(CommandHandler("start", start, filters=AuthorizedUserFilter())) - application.add_handler(CommandHandler("help", help_command, filters=AuthorizedUserFilter())) - application.add_handler(CommandHandler("new", newchat_command, filters=AuthorizedUserFilter())) + application.add_handler(CommandHandler("start", start, filters=AuthFilter)) + application.add_handler(CommandHandler("help", help_command, filters=AuthFilter)) + application.add_handler(CommandHandler("new", newchat_command, filters=AuthFilter)) # Any text message is sent to LLM to generate a response - application.add_handler( - MessageHandler( AuthorizedUserFilter() & ~filters.COMMAND & filters.TEXT, handle_message) - ) + application.add_handler(MessageHandler(MessageFilter, handle_message)) # Any image is sent to LLM to generate a response - application.add_handler( - MessageHandler( AuthorizedUserFilter() & ~filters.COMMAND & filters.PHOTO, handle_image) - ) + application.add_handler(MessageHandler(PhotoFilter, handle_image)) # Run the bot until the user presses Ctrl-C application.run_polling(allowed_updates=Update.ALL_TYPES) diff --git a/gemini_pro_bot/filters.py b/gemini_pro_bot/filters.py index f7ea0c7..d271a82 100644 --- a/gemini_pro_bot/filters.py +++ b/gemini_pro_bot/filters.py @@ -1,13 +1,25 @@ import os from telegram import Update -from telegram.ext.filters import UpdateFilter +from telegram.ext.filters import UpdateFilter, COMMAND, TEXT, PHOTO from dotenv import load_dotenv load_dotenv() -AUTHORIZED_USERS = [i.strip() for i in os.getenv("AUTHORIZED_USERS", "").split(",") if i.strip()] +_AUTHORIZED_USERS = [ + i.strip() for i in os.getenv("AUTHORIZED_USERS", "").split(",") if i.strip() +] + + class AuthorizedUserFilter(UpdateFilter): def filter(self, update: Update): - if not AUTHORIZED_USERS: + if not _AUTHORIZED_USERS: return True - return update.message.from_user.username in AUTHORIZED_USERS or str(update.message.from_user.id) in AUTHORIZED_USERS \ No newline at end of file + return ( + update.message.from_user.username in _AUTHORIZED_USERS + or str(update.message.from_user.id) in _AUTHORIZED_USERS + ) + + +AuthFilter = AuthorizedUserFilter() +MessageFilter = AuthFilter & ~COMMAND & TEXT +PhotoFilter = AuthFilter & ~COMMAND & PHOTO