From 6ffae054c0126a1e2f3ab3af22eb2c4a1582628d Mon Sep 17 00:00:00 2001 From: Ilya Bezrukov Date: Sun, 21 Jul 2024 01:58:18 +0300 Subject: [PATCH] Fix arguments middleware and add state storage --- mybot/bot.py | 14 +++++++++++++- mybot/config.py | 7 +++++++ mybot/filters/__init__.py | 0 mybot/handlers/basic.py | 10 +++++----- mybot/i18n.py | 8 ++++---- mybot/middlewares/__init__.py | 2 +- mybot/middlewares/arguments.py | 24 ++++++------------------ mybot/states.py | 1 + 8 files changed, 37 insertions(+), 29 deletions(-) create mode 100644 mybot/filters/__init__.py create mode 100644 mybot/states.py diff --git a/mybot/bot.py b/mybot/bot.py index 37d609f..d4f7b51 100644 --- a/mybot/bot.py +++ b/mybot/bot.py @@ -1,15 +1,27 @@ from telebot import TeleBot +from telebot.storage import StateMemoryStorage, StateRedisStorage from .handlers import register_handlers from .middlewares import setup_middlewares def create_bot(config, logger, i18n): + if config.SS_TYPE == "memory": + state_storage = StateMemoryStorage() + elif config.SS_TYPE == "redis": + state_storage = StateRedisStorage(config.SS_REDIS_HOST, + config.SS_REDIS_PORT, + config.SS_REDIS_DB, + config.SS_REDIS_PASS) + else: + raise RuntimeWarning(f"Unknown state storage type: '{config.SS_TYPE}'") + bot = TeleBot(config.BOT_TOKEN, parse_mode=config.PARSE_MODE, skip_pending=config.DROP_PENDING, num_threads=config.NUM_THREADS, - use_class_middlewares=True) + use_class_middlewares=True, + state_storage=state_storage) logger.debug("Setting up middlewares") setup_middlewares(bot, logger, i18n) diff --git a/mybot/config.py b/mybot/config.py index d71acaa..2433c74 100644 --- a/mybot/config.py +++ b/mybot/config.py @@ -18,6 +18,13 @@ class Config: I18N_PATH: str = os.getenv("I18N_PATH", "i18n.yaml") I18N_LANG: str = os.getenv("I18N_LANG", "en") + # state storage + SS_TYPE: str = os.getenv("SS_TYPE", "memory").lower() + SS_REDIS_HOST: str = os.getenv("SS_REDIS_HOST") + SS_REDIS_PORT: int = int(os.getenv("SS_REDIS_PORT"), 6379) + SS_REDIS_DB: int = int(os.getenv("SS_REDIS_DB", 0)) + SS_REDIS_PASS: str = os.getenv("SS_REDIS_PASS") + def __init__(self): if not self.BOT_TOKEN: raise RuntimeError("Missing BOT_TOKEN") diff --git a/mybot/filters/__init__.py b/mybot/filters/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mybot/handlers/basic.py b/mybot/handlers/basic.py index 0198d5f..9344143 100644 --- a/mybot/handlers/basic.py +++ b/mybot/handlers/basic.py @@ -1,15 +1,15 @@ from telebot import TeleBot -from telebot import types as tt +from telebot.types import Message -def start(message: tt.Message, bot: TeleBot, logger, t): +def start(message: Message, bot: TeleBot, t, **kwargs): bot.send_message(message.chat.id, t("start")) -def help_(message: tt.Message, bot: TeleBot, logger, t): +def help_(message, bot, t, **kwargs): bot.send_message(message.chat.id, t("help")) def register_handlers(bot: TeleBot): - bot.register_message_handler(start, commands=["start"]) - bot.register_message_handler(help_, commands=["help"]) + bot.register_message_handler(start, commands=["start"], pass_bot=True) + bot.register_message_handler(help_, commands=["help"], pass_bot=True) diff --git a/mybot/i18n.py b/mybot/i18n.py index 38d08cf..226b79f 100644 --- a/mybot/i18n.py +++ b/mybot/i18n.py @@ -42,7 +42,7 @@ class I18N: def __call__(self, phrase: str, lang: Optional[str] = None, *args, **kwargs): return self.get(phrase, lang).format(*args, **kwargs) - def custom_call(self, lang: Optional[str] = None, *args, **kwargs): - def call(phrase: str, ilang: Optional[str] = None, *iargs, **ikwargs): - return self(phrase, lang or ilang, *args, *iargs, **kwargs, **ikwargs) - return call + def customized_call(self, preset_lang: Optional[str] = None, **preset_kwargs): + def wrapper(phrase: str, lang: Optional[str] = None, *args, **kwargs): + return self(phrase, lang or preset_lang, *args, **preset_kwargs, **kwargs) + return wrapper diff --git a/mybot/middlewares/__init__.py b/mybot/middlewares/__init__.py index 23b0449..79fdc02 100644 --- a/mybot/middlewares/__init__.py +++ b/mybot/middlewares/__init__.py @@ -6,4 +6,4 @@ from .arguments import ExtraArguments def setup_middlewares(bot: TeleBot, logger: logging.Logger, i18n): - bot.setup_middleware(ExtraArguments(bot, logger, i18n)) + bot.setup_middleware(ExtraArguments(logger, i18n)) diff --git a/mybot/middlewares/arguments.py b/mybot/middlewares/arguments.py index f712b97..5603e66 100644 --- a/mybot/middlewares/arguments.py +++ b/mybot/middlewares/arguments.py @@ -1,34 +1,22 @@ import logging -from functools import singledispatch -from telebot import TeleBot from telebot.handler_backends import BaseMiddleware -from telebot import types as tt +from telebot.types import Message, CallbackQuery class ExtraArguments(BaseMiddleware): - def __init__(self, bot: TeleBot, logger: logging.Logger, i18n): + def __init__(self, logger: logging.Logger, i18n): super().__init__() - self.bot = bot self.logger = logger self.i18n = i18n self.update_types = ["message", "callback_query"] - @singledispatch def pre_process(self, obj, data: dict): - raise TypeError() - - @pre_process.register - def _(self, message: tt.Message, data: dict): - data["bot"] = self.bot data["logger"] = self.logger - data["t"] = self.i18n.custom_call(message=message) - - @pre_process.register - def _(self, callback: tt.CallbackQuery, data: dict): - data["bot"] = self.bot - data["logger"] = self.logger - data["t"] = self.i18n.custom_call(callback=callback) + if isinstance(obj, Message): + data["t"] = self.i18n.customized_call(message=obj) + elif isinstance(obj, CallbackQuery): + data["t"] = self.i18n.customized_call(callback=obj) def post_process(self, message, data: dict, exception: BaseException): pass diff --git a/mybot/states.py b/mybot/states.py new file mode 100644 index 0000000..03b7d55 --- /dev/null +++ b/mybot/states.py @@ -0,0 +1 @@ +# states will be defined here