Fix arguments middleware and add state storage

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)

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")

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)

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

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))

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

mybot/states.py

@@ -0,0 +1 @@
+# states will be defined here