From 273c0bf2dadb583ef9142f8f454b5e37be2ba366 Mon Sep 17 00:00:00 2001 From: Ilya Bezrukov Date: Fri, 26 Jul 2024 02:58:19 +0300 Subject: [PATCH] Refactored config --- mybot/__init__.py | 31 +++++++--- mybot/bot.py | 33 ++--------- mybot/config.py | 104 ++++++++++++++++++++++++--------- mybot/database/__init__.py | 6 +- mybot/filters/__init__.py | 5 +- mybot/i18n.py | 19 +++--- mybot/middlewares/__init__.py | 7 +-- mybot/middlewares/arguments.py | 6 +- mybot/states.py | 17 ++++++ requirements.txt | 1 + 10 files changed, 144 insertions(+), 85 deletions(-) diff --git a/mybot/__init__.py b/mybot/__init__.py index 971c3bf..791b780 100644 --- a/mybot/__init__.py +++ b/mybot/__init__.py @@ -1,12 +1,29 @@ -from .config import Config +from .config import Config, load_config from .logger import create_logger -from .bot import create_bot +from .bot import get_bot from .i18n import I18N +from .database import get_engine +from .states import get_state_storage +from .middlewares import setup_middlewares +from .filters import add_custom_filters + + +def create_bot(config: Config, i18n: I18N): + state_storage = get_state_storage(config.states) + bot = get_bot(config.bot, state_storage) + setup_middlewares(bot, i18n) + add_custom_filters(bot, config) + return bot def main(): - config = Config() - logger = create_logger("mybot", config.LOG_LEVEL) - i18n = I18N(logger, config.I18N_PATH, config.I18N_LANG, config.I18N_FALLBACK_LANG) - bot = create_bot(config, logger, i18n) - bot.infinity_polling(config.TIMEOUT, config.DROP_PENDING, config.POLLING_TIMEOUT) + config = load_config() + # logger = create_logger("mybot", config.log_level) + i18n = I18N(config.i18n) + # engine = get_engine(config.database) + bot = create_bot(config, i18n) + bot.infinity_polling( + timeout=config.bot.timeout, + long_polling_timeout=config.bot.polling_timeout, + skip_pending=config.bot.skip_pending + ) diff --git a/mybot/bot.py b/mybot/bot.py index 478c4f6..e5c4a65 100644 --- a/mybot/bot.py +++ b/mybot/bot.py @@ -1,36 +1,15 @@ from telebot import TeleBot -from telebot.storage import StateMemoryStorage, StateRedisStorage from .handlers import register_handlers -from .middlewares import setup_middlewares -from .filters import add_custom_filters +from .config import BotConfig -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, +def get_bot(config: BotConfig, state_storage): + bot = TeleBot(config.token, + parse_mode=config.parse_mode, + skip_pending=config.skip_pending, + num_threads=config.num_threads, use_class_middlewares=True, state_storage=state_storage) - - logger.debug("Setting up middlewares") - setup_middlewares(bot, logger, i18n) - - logger.debug("Registering handlers") register_handlers(bot) - - logger.debug("Adding custom filters") - add_custom_filters(bot, config) - return bot diff --git a/mybot/config.py b/mybot/config.py index f79dfd8..74c6588 100644 --- a/mybot/config.py +++ b/mybot/config.py @@ -1,34 +1,86 @@ import os +from dataclasses import dataclass +@dataclass +class BotConfig: + token: str + skip_pending: bool + timeout: int + polling_timeout: int + num_threads: int + parse_mode: str + + @classmethod + def from_env(cls): + return cls(os.getenv("BOT_TOKEN"), + bool(int(os.getenv("BOT_SKIP_PENDING", True))), + int(os.getenv("BOT_TIMEOUT", 20)), + int(os.getenv("BOT_POLLING_TIMEOUT", 20)), + int(os.getenv("BOT_NUM_THREADS", 2)), + os.getenv("BOT_PARSE_MODE", "html")) + + +@dataclass +class I18NConfig: + path: str + lang: str + fallback_lang: str + + @classmethod + def from_env(cls): + return cls(os.getenv("I18N_PATH", "i18n.yaml"), + os.getenv("I18N_LANG", "en"), + os.getenv("I18N_FALLBACK_LANG", "en")) + + +@dataclass +class StateStorageConfig: + type: str + redis_host: str + redis_port: int + redis_db: int + redis_pass: str + + @classmethod + def from_env(cls): + return cls(os.getenv("SS_TYPE"), + os.getenv("SS_REDIS_HOST"), + int(os.getenv("SS_REDIS_PORT", 6379)), + int(os.getenv("SS_REDIS_DB", 0)), + os.getenv("SS_REDIS_PASS")) + + +@dataclass +class DatabaseConfig: + url: str + + @classmethod + def from_env(cls): + return cls(os.getenv("DATABASE_URL", "sqlite:///bot.db")) + + +@dataclass class Config: - # bot setup - BOT_TOKEN: str = os.getenv("BOT_TOKEN") - OWNER_ID: int = int(os.getenv("OWNER_ID", 1)) - LOG_LEVEL: str = os.getenv("LOG_LEVEL", "INFO") + bot: BotConfig + i18n: I18NConfig + states: StateStorageConfig + database: DatabaseConfig - # bot behaviour - DROP_PENDING: bool = bool(int(os.getenv("DROP_PENDING", True))) - TIMEOUT: int = int(os.getenv("TIMEOUT", 20)) - POLLING_TIMEOUT: int = int(os.getenv("POLLING_TIMEOUT", 20)) - NUM_THREADS: int = int(os.getenv("NUM_THREADS", 2)) - PARSE_MODE: str = os.getenv("PARSE_MODE", "html") + log_level: str + owner_id: int - # i18n - I18N_PATH: str = os.getenv("I18N_PATH", "i18n.yaml") - I18N_LANG: str = os.getenv("I18N_LANG", "en") - I18N_FALLBACK_LANG: str = os.getenv("I18N_FALLBACK_LANG", "en") + @classmethod + def from_env(cls): + return cls( + bot=BotConfig.from_env(), + i18n=I18NConfig.from_env(), + states=StateStorageConfig.from_env(), + database=DatabaseConfig.from_env(), + log_level=os.getenv("LOG_LEVEL", "INFO"), + owner_id=int(os.getenv("OWNER_ID", 1)), + ) - # 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") - # database - DB_URL: str = os.getenv("DB_URL", "sqlite:///bot.db") - - def __init__(self): - if not self.BOT_TOKEN: - raise RuntimeError("Missing BOT_TOKEN") +def load_config() -> Config: + return Config.from_env() diff --git a/mybot/database/__init__.py b/mybot/database/__init__.py index e3442d8..b0d6015 100644 --- a/mybot/database/__init__.py +++ b/mybot/database/__init__.py @@ -1,11 +1,11 @@ from sqlalchemy import create_engine from sqlalchemy.orm import DeclarativeBase -from ..config import Config +from ..config import DatabaseConfig -def get_engine(config: Config): - engine = create_engine(config.DB_URL, +def get_engine(config: DatabaseConfig): + engine = create_engine(config.url, pool_recycle=3600, pool_pre_ping=True) return engine diff --git a/mybot/filters/__init__.py b/mybot/filters/__init__.py index 377fad7..f7eed56 100644 --- a/mybot/filters/__init__.py +++ b/mybot/filters/__init__.py @@ -2,8 +2,9 @@ from telebot import TeleBot from telebot.custom_filters import StateFilter from .isowner import IsOwnerFilter +from ..config import Config -def add_custom_filters(bot: TeleBot, config): +def add_custom_filters(bot: TeleBot, config: Config): bot.add_custom_filter(StateFilter(bot)) - bot.add_custom_filter(IsOwnerFilter(config.OWNER_ID)) + bot.add_custom_filter(IsOwnerFilter(config.owner_id)) diff --git a/mybot/i18n.py b/mybot/i18n.py index 9684ec5..954ef2e 100644 --- a/mybot/i18n.py +++ b/mybot/i18n.py @@ -1,18 +1,15 @@ -import logging from typing import Optional from yaml import safe_load +from .config import I18NConfig + class I18N: - def __init__(self, logger: logging.Logger, - path: str = "i18n.yaml", - lang: Optional[str] = None, - fallback_lang: str = "en"): - self.logger = logger - self._path = path - self._fallback_lang = fallback_lang - self._lang = lang or self.fallback_lang + def __init__(self, config: I18NConfig): + self._path = config.path + self._fallback_lang = config.fallback_lang + self._lang = config.lang or self.fallback_lang self._dict = dict() self.load() @@ -44,12 +41,12 @@ class I18N: def get(self, phrase: str, lang: Optional[str] = None): lang = lang or self.lang if lang not in self._dict: - self.logger.warning(f"Language '{lang}' not found in i18n, using fallback") + # self.logger.warning(f"Language '{lang}' not found in i18n, using fallback") lang = self.fallback_lang lang_dict = self._dict.get(lang) result = lang_dict.get(phrase) if result is None: - self.logger.error(f"Phrase '{phrase}' not found in language '{lang}'") + # self.logger.error(f"Phrase '{phrase}' not found in language '{lang}'") result = f"" return result diff --git a/mybot/middlewares/__init__.py b/mybot/middlewares/__init__.py index 79fdc02..26d0342 100644 --- a/mybot/middlewares/__init__.py +++ b/mybot/middlewares/__init__.py @@ -1,9 +1,8 @@ -import logging - from telebot import TeleBot from .arguments import ExtraArguments +from ..i18n import I18N -def setup_middlewares(bot: TeleBot, logger: logging.Logger, i18n): - bot.setup_middleware(ExtraArguments(logger, i18n)) +def setup_middlewares(bot: TeleBot, i18n: I18N): + bot.setup_middleware(ExtraArguments(i18n)) diff --git a/mybot/middlewares/arguments.py b/mybot/middlewares/arguments.py index 5603e66..1b9c07c 100644 --- a/mybot/middlewares/arguments.py +++ b/mybot/middlewares/arguments.py @@ -1,18 +1,14 @@ -import logging - from telebot.handler_backends import BaseMiddleware from telebot.types import Message, CallbackQuery class ExtraArguments(BaseMiddleware): - def __init__(self, logger: logging.Logger, i18n): + def __init__(self, i18n): super().__init__() - self.logger = logger self.i18n = i18n self.update_types = ["message", "callback_query"] def pre_process(self, obj, data: dict): - data["logger"] = self.logger if isinstance(obj, Message): data["t"] = self.i18n.customized_call(message=obj) elif isinstance(obj, CallbackQuery): diff --git a/mybot/states.py b/mybot/states.py index 03b7d55..d0ead71 100644 --- a/mybot/states.py +++ b/mybot/states.py @@ -1 +1,18 @@ +from telebot.storage import StateMemoryStorage, StateRedisStorage + +from .config import StateStorageConfig + + +def get_state_storage(config: StateStorageConfig): + if config.type == "memory": + state_storage = StateMemoryStorage() + elif config.type == "redis": + state_storage = StateRedisStorage(config.redis_host, config.redis_port, + config.redis_db, config.redis_pass) + else: + raise RuntimeWarning(f"Unknown state storage type: '{config.type}'") + return state_storage + + # states will be defined here + diff --git a/requirements.txt b/requirements.txt index 8e2d4db..6bda7df 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ pytelegrambotapi +environs pyyaml sqlalchemy alembic