From 02fe97e921708a83a37020a5c100643f362030a2 Mon Sep 17 00:00:00 2001 From: Alexander Kalinovsky Date: Mon, 10 Feb 2025 20:24:29 +0100 Subject: [PATCH] upd bot init --- lifespan.py | 74 ---------------------------------- main.py | 112 +++++++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 110 insertions(+), 76 deletions(-) delete mode 100644 lifespan.py diff --git a/lifespan.py b/lifespan.py deleted file mode 100644 index 1f73d96..0000000 --- a/lifespan.py +++ /dev/null @@ -1,74 +0,0 @@ -from aiogram.types import BotCommand -from contextlib import asynccontextmanager -from .main import QBotApp -from logging import getLogger - -logger = getLogger(__name__) - - -@asynccontextmanager -async def default_lifespan(app: QBotApp): - logger.debug("starting qbot app") - - if app.config.USE_NGROK: - try: - from pyngrok import ngrok - from pyngrok.conf import PyngrokConfig - - except ImportError: - logger.error("pyngrok is not installed") - raise - - tunnel = ngrok.connect( - app.config.API_PORT, - pyngrok_config=PyngrokConfig(auth_token=app.config.NGROK_AUTH_TOKEN), - ) - app.config.NGROK_URL = tunnel.public_url - - commands_captions = dict[str, list[tuple[str, str]]]() - - for command_name, command in app.bot_commands.items(): - if command.show_in_bot_commands: - if isinstance(command.caption, str) or command.caption is None: - if "default" not in commands_captions: - commands_captions["default"] = [] - commands_captions["default"].append( - (command_name, command.caption or command_name) - ) - else: - for locale, description in command.caption.items(): - if locale not in commands_captions: - commands_captions[locale] = [] - commands_captions[locale].append((command_name, description)) - - for locale, commands in commands_captions.items(): - await app.bot.set_my_commands( - [ - BotCommand(command=command[0], description=command[1]) - for command in commands - ], - language_code=None if locale == "default" else locale, - ) - - await app.bot.set_webhook( - url=f"{app.config.API_URL}/api/telegram/webhook", - drop_pending_updates=True, - allowed_updates=["message", "callback_query", "pre_checkout_query"], - secret_token=app.bot_auth_token, - ) - - logger.info("qbot app started") - - if app.lifespan: - async with app.lifespan(app): - yield - else: - yield - - logger.info("stopping qbot app") - - await app.bot.delete_webhook() - if app.config.USE_NGROK: - ngrok.disconnect(app.config.NGROK_URL) - ngrok.kill() - logger.info("qbot app stopped") diff --git a/main.py b/main.py index e12ce1c..ec49454 100644 --- a/main.py +++ b/main.py @@ -1,13 +1,15 @@ +from contextlib import asynccontextmanager from typing import Annotated, Callable, Any from typing_extensions import Doc from aiogram import Bot, Dispatcher from aiogram.client.default import DefaultBotProperties -from aiogram.types import Message +from aiogram.types import Message, BotCommand as AiogramBotCommand from aiogram.utils.callback_answer import CallbackAnswerMiddleware from aiogram.utils.i18n import I18n from fastapi import FastAPI from fastapi.applications import Lifespan, AppType from secrets import token_hex +from logging import getLogger from .config import Config from .fsm.db_storage import DbStorage @@ -18,6 +20,38 @@ from .model.descriptors import BotCommand from .router import Router +logger = getLogger(__name__) + +@asynccontextmanager +async def default_lifespan(app: "QBotApp"): + logger.debug("starting qbot app") + + if app.lifespan_bot_init: + + if app.config.USE_NGROK: + app.ngrok_init() + + await app.bot_init() + + logger.info("qbot app started") + + if app.lifespan: + async with app.lifespan(app): + yield + else: + yield + + logger.info("stopping qbot app") + + if app.lifespan_bot_init: + await app.bot_close() + + if app.config.USE_NGROK: + app.ngrok_stop() + + logger.info("qbot app stopped") + + class QBotApp(FastAPI): """ Main class for the QBot application @@ -50,6 +84,8 @@ class QBotApp(FastAPI): | None ) = None, lifespan: Lifespan[AppType] | None = None, + lifespan_bot_init: bool = True, + allowed_updates: list[str] | None = None, *args, **kwargs, ): @@ -61,6 +97,8 @@ class QBotApp(FastAPI): user_class = DefaultUser + self.allowed_updates = allowed_updates or ["message", "callback_query"] + self.user_class = user_class self.entity_metadata: EntityMetadata = user_class.entity_metadata self.config = config @@ -95,7 +133,7 @@ class QBotApp(FastAPI): self.start_handler = bot_start self.bot_commands = dict[str, BotCommand]() - from .lifespan import default_lifespan + self.lifespan_bot_init = lifespan_bot_init super().__init__(lifespan=default_lifespan, *args, **kwargs) @@ -106,7 +144,77 @@ class QBotApp(FastAPI): self.root_router._commands = self.bot_commands self.command = self.root_router.command + def register_routers(self, *routers: Router): for router in routers: for command_name, command in router._commands.items(): self.bot_commands[command_name] = command + + + def ngrok_init(self): + try: + from pyngrok import ngrok + from pyngrok.conf import PyngrokConfig + + except ImportError: + logger.error("pyngrok is not installed") + raise + + tunnel = ngrok.connect( + self.config.API_PORT, + pyngrok_config=PyngrokConfig(auth_token=self.config.NGROK_AUTH_TOKEN), + ) + self.config.NGROK_URL = tunnel.public_url + + + def ngrok_stop(self): + try: + from pyngrok import ngrok + + except ImportError: + logger.error("pyngrok is not installed") + raise + + ngrok.disconnect(self.config.NGROK_URL) + ngrok.kill() + + + + async def bot_init(self): + + commands_captions = dict[str, list[tuple[str, str]]]() + + for command_name, command in self.bot_commands.items(): + if command.show_in_bot_commands: + if isinstance(command.caption, str) or command.caption is None: + if "default" not in commands_captions: + commands_captions["default"] = [] + commands_captions["default"].append( + (command_name, command.caption or command_name) + ) + else: + for locale, description in command.caption.items(): + if locale not in commands_captions: + commands_captions[locale] = [] + commands_captions[locale].append((command_name, description)) + + for locale, commands in commands_captions.items(): + await self.bot.set_my_commands( + [ + AiogramBotCommand(command=command[0], description=command[1]) + for command in commands + ], + language_code=None if locale == "default" else locale, + ) + + await self.bot.set_webhook( + url=f"{self.config.API_URL}/api/telegram/webhook", + drop_pending_updates=True, + allowed_updates=self.allowed_updates, + secret_token=self.bot_auth_token, + ) + + + async def bot_close(self): + await self.bot.delete_webhook() + \ No newline at end of file