From ea58ee41373576cfb2394a8654af2aab2aee0bc5 Mon Sep 17 00:00:00 2001 From: Forevka <32968153+Forevka@users.noreply.github.com> Date: Sat, 3 Jul 2021 12:02:39 +0000 Subject: [PATCH] provide allowed_updates in polling mode --- aiogram/dispatcher/dispatcher.py | 13 ++++-- aiogram/utils/handlers_in_use.py | 19 +++++++++ examples/specify_updates.py | 72 ++++++++++++++++++++++++++++++++ 3 files changed, 101 insertions(+), 3 deletions(-) create mode 100644 aiogram/utils/handlers_in_use.py create mode 100644 examples/specify_updates.py diff --git a/aiogram/dispatcher/dispatcher.py b/aiogram/dispatcher/dispatcher.py index 47c67096..48982177 100644 --- a/aiogram/dispatcher/dispatcher.py +++ b/aiogram/dispatcher/dispatcher.py @@ -4,7 +4,7 @@ import asyncio import contextvars import warnings from asyncio import CancelledError, Future, Lock -from typing import Any, AsyncGenerator, Dict, Optional, Union +from typing import Any, AsyncGenerator, Dict, List, Optional, Union from .. import loggers from ..client.bot import Bot @@ -130,6 +130,7 @@ class Dispatcher(Router): bot: Bot, polling_timeout: int = 30, backoff_config: BackoffConfig = DEFAULT_BACKOFF_CONFIG, + allowed_updates: Optional[List[str]] = None, ) -> AsyncGenerator[Update, None]: """ Endless updates reader with correctly handling any server-side or connection errors. @@ -137,7 +138,7 @@ class Dispatcher(Router): So you may not worry that the polling will stop working. """ backoff = Backoff(config=backoff_config) - get_updates = GetUpdates(timeout=polling_timeout) + get_updates = GetUpdates(timeout=polling_timeout, allowed_updates=allowed_updates) kwargs = {} if bot.session.timeout: # Request timeout can be lower than session timeout ant that's OK. @@ -297,6 +298,7 @@ class Dispatcher(Router): polling_timeout: int = 30, handle_as_tasks: bool = True, backoff_config: BackoffConfig = DEFAULT_BACKOFF_CONFIG, + allowed_updates: Optional[List[str]] = None, **kwargs: Any, ) -> None: """ @@ -307,7 +309,7 @@ class Dispatcher(Router): :return: """ async for update in self._listen_updates( - bot, polling_timeout=polling_timeout, backoff_config=backoff_config + bot, polling_timeout=polling_timeout, backoff_config=backoff_config, allowed_updates=allowed_updates, ): handle_update = self._process_update(bot=bot, update=update, **kwargs) if handle_as_tasks: @@ -397,6 +399,7 @@ class Dispatcher(Router): polling_timeout: int = 10, handle_as_tasks: bool = True, backoff_config: BackoffConfig = DEFAULT_BACKOFF_CONFIG, + allowed_updates: Optional[List[str]] = None, **kwargs: Any, ) -> None: """ @@ -427,6 +430,7 @@ class Dispatcher(Router): handle_as_tasks=handle_as_tasks, polling_timeout=polling_timeout, backoff_config=backoff_config, + allowed_updates=allowed_updates, **kwargs, ) ) @@ -443,6 +447,7 @@ class Dispatcher(Router): polling_timeout: int = 30, handle_as_tasks: bool = True, backoff_config: BackoffConfig = DEFAULT_BACKOFF_CONFIG, + allowed_updates: Optional[List[str]] = None, **kwargs: Any, ) -> None: """ @@ -463,8 +468,10 @@ class Dispatcher(Router): polling_timeout=polling_timeout, handle_as_tasks=handle_as_tasks, backoff_config=backoff_config, + allowed_updates=allowed_updates, ) ) except (KeyboardInterrupt, SystemExit): # pragma: no cover # Allow to graceful shutdown pass + \ No newline at end of file diff --git a/aiogram/utils/handlers_in_use.py b/aiogram/utils/handlers_in_use.py new file mode 100644 index 00000000..925870f5 --- /dev/null +++ b/aiogram/utils/handlers_in_use.py @@ -0,0 +1,19 @@ + +from itertools import chain +from typing import List +from aiogram.dispatcher.dispatcher import Dispatcher + +AIOGRAM_INTERNAL_HANDLERS = ['update', 'error',] + +def get_handlers_in_use(dispatcher: Dispatcher, handlers_to_skip: List[str] = AIOGRAM_INTERNAL_HANDLERS) -> List[str]: + handlers_in_use = [] + + for router in [dispatcher.sub_routers, dispatcher]: + if (isinstance(router, list)): + if (router): + handlers_in_use.extend(chain(*list(map(get_handlers_in_use, router)))) + else: + for update_name, observer in router.observers.items(): + if (observer.handlers and update_name not in [*handlers_to_skip, *handlers_in_use]): handlers_in_use.append(update_name) + + return handlers_in_use diff --git a/examples/specify_updates.py b/examples/specify_updates.py new file mode 100644 index 00000000..17c49072 --- /dev/null +++ b/examples/specify_updates.py @@ -0,0 +1,72 @@ +from aiogram.types.inline_keyboard_button import InlineKeyboardButton +from aiogram.types.inline_keyboard_markup import InlineKeyboardMarkup +from aiogram.dispatcher.router import Router +from aiogram.utils.handlers_in_use import get_handlers_in_use +import logging + +from aiogram import Bot, Dispatcher +from aiogram.types import Message, ChatMemberUpdated, CallbackQuery + +TOKEN = "666922879:AAEWkOwKYH-Sz7pBm9fLtXDlDV1fSGiNbwo" +dp = Dispatcher() + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + +@dp.message(commands={"start"}) +async def command_start_handler(message: Message) -> None: + """ + This handler receive messages with `/start` command + """ + + await message.answer(f"Hello, {message.from_user.full_name}!", reply_markup=InlineKeyboardMarkup(inline_keyboard=[[InlineKeyboardButton(text="Tap me, bro", callback_data="*")]])) + + +@dp.chat_member() +async def chat_member_update(chat_member: ChatMemberUpdated, bot: Bot) -> None: + await bot.send_message(chat_member.chat.id, f"Member {chat_member.from_user.id} was changed from {chat_member.old_chat_member.is_chat_member} to {chat_member.new_chat_member.is_chat_member}") + +# this router will use only callback_query updates +sub_router = Router() + +@sub_router.callback_query() +async def callback_tap_me(callback_query: CallbackQuery) -> None: + await callback_query.answer("Yeah good, now i'm fine") + + +# this router will use only edited_message updates +sub_sub_router = Router() + +@sub_sub_router.edited_message() +async def callback_tap_me(edited_message: Message) -> None: + await edited_message.reply("Message was edited, big brother watch you") + + +# this router will use only my_chat_member updates +deep_dark_router = Router() + +@deep_dark_router.my_chat_member() +async def my_chat_member_change(chat_member: ChatMemberUpdated, bot: Bot) -> None: + await bot.send_message(chat_member.chat.id, f"Member was changed from {chat_member.old_chat_member.is_chat_member} to {chat_member.new_chat_member.is_chat_member}") + + +def main() -> None: + # Initialize Bot instance with an default parse mode which will be passed to all API calls + bot = Bot(TOKEN, parse_mode="HTML") + + sub_router.include_router(deep_dark_router) + + dp.include_router(sub_router) + dp.include_router(sub_sub_router) + + useful_updates = get_handlers_in_use(dp) + + # And the run events dispatching + dp.run_polling( + bot, + allowed_updates=useful_updates + ) + + +if __name__ == "__main__": + main()