From 1e242f22634f8504e456d161d083757a1d075ce8 Mon Sep 17 00:00:00 2001 From: bedilbek Date: Wed, 8 Apr 2020 23:13:19 +0500 Subject: [PATCH] Add Middleware support --- telebot/__init__.py | 74 +++++++++++++++++++++++++++++++++++++++++++ tests/test_telebot.py | 49 ++++++++++++++++++++++++++++ 2 files changed, 123 insertions(+) diff --git a/telebot/__init__.py b/telebot/__init__.py index b7b5d79..a143479 100644 --- a/telebot/__init__.py +++ b/telebot/__init__.py @@ -169,6 +169,21 @@ class TeleBot: self.pre_checkout_query_handlers = [] self.poll_handlers = [] + self.typed_middleware_handlers = { + 'message': [], + 'edited_message': [], + 'channel_post': [], + 'edited_channel_post': [], + 'inline_query': [], + 'chosen_inline_result': [], + 'callback_query': [], + 'shipping_query': [], + 'pre_checkout_query': [], + 'poll': [], + } + + self.default_middleware_handlers = [] + self.threaded = threaded if self.threaded: self.worker_pool = util.ThreadPool(num_threads=num_threads) @@ -293,6 +308,9 @@ class TeleBot: new_polls = [] for update in updates: + + self.process_middlewares(update) + if update.update_id > self.last_update_id: self.last_update_id = update.update_id if update.message: @@ -371,6 +389,16 @@ class TeleBot: def process_new_poll(self, polls): self._notify_command_handlers(self.poll_handlers, polls) + def process_middlewares(self, update): + for update_type, middlewares in self.typed_middleware_handlers.items(): + if hasattr(update, update_type): + for typed_middleware_handler in middlewares: + typed_middleware_handler(self, getattr(update, update_type)) + + if len(self.default_middleware_handlers) > 0: + for default_middleware_handler in self.default_middleware_handlers: + default_middleware_handler(self, update) + def __notify_update(self, new_messages): for listener in self.update_listener: self._exec_task(listener, new_messages) @@ -1497,6 +1525,52 @@ class TeleBot: 'filters' : filters } + def middleware_handler(self, update_types=None): + """ + Middleware handler decorator. + + This decorator can be used to decorate functions that must be handled as middlewares before entering any other + message handlers + But, be careful and check type of the update inside the handler if more than one update_type is given + + Example: + + bot = TeleBot('TOKEN') + + # Print post message text before entering to any post_channel handlers + @bot.middleware_handler(update_types=['channel_post', 'edited_channel_post']) + def print_channel_post_text(bot_instance, channel_post): + print(channel_post.text) + + # Print update id before entering to any handlers + @bot.middleware_handler() + def print_channel_post_text(bot_instance, update): + print(update.update_id) + + :param update_types: Optional list of update types that can be passed into the middleware handler. + + """ + + def decorator(handler): + self.add_middleware_handler(handler, update_types) + + return handler + + return decorator + + def add_middleware_handler(self, handler, update_types=None): + """ + Add middleware handler + :param handler: + :param update_types: + :return: + """ + if update_types: + for update_type in update_types: + self.typed_middleware_handlers[update_type].append(handler) + else: + self.default_middleware_handlers.append(handler) + def message_handler(self, commands=None, regexp=None, func=None, content_types=None, **kwargs): """ Message handler decorator. diff --git a/tests/test_telebot.py b/tests/test_telebot.py index 908b141..3157bdd 100644 --- a/tests/test_telebot.py +++ b/tests/test_telebot.py @@ -408,6 +408,23 @@ class TestTeleBot: chat = types.User(11, False, 'test') return types.Message(1, None, None, chat, 'text', params, "") + @staticmethod + def create_message_update(text): + params = {'text': text} + chat = types.User(11, False, 'test') + message = types.Message(1, None, None, chat, 'text', params, "") + edited_message = None + channel_post = None + edited_channel_post = None + inline_query = None + chosen_inline_result = None + callback_query = None + shipping_query = None + pre_checkout_query = None + poll = None + return types.Update(-1001234038283, message, edited_message, channel_post, edited_channel_post, inline_query, + chosen_inline_result, callback_query, shipping_query, pre_checkout_query, poll) + def test_is_string_unicode(self): s1 = u'string' assert util.is_string(s1) @@ -489,3 +506,35 @@ class TestTeleBot: tb = telebot.TeleBot(TOKEN) ret_msg = tb.send_document(CHAT_ID, file_data, caption='_italic_', parse_mode='Markdown') assert ret_msg.caption_entities[0].type == 'italic' + + def test_typed_middleware_handler(self): + tb = telebot.TeleBot('') + update = self.create_message_update('/help') + + @tb.middleware_handler(update_types=['message']) + def middleware(tb_instance, message): + message.text = 'got' + + @tb.message_handler(func=lambda m: m.text == 'got') + def command_handler(message): + message.text = message.text + message.text + + tb.process_new_updates([update]) + time.sleep(1) + assert update.message.text == 'got' * 2 + + def test_default_middleware_handler(self): + tb = telebot.TeleBot('') + update = self.create_message_update('/help') + + @tb.middleware_handler() + def middleware(tb_instance, update): + update.message.text = 'got' + + @tb.message_handler(func=lambda m: m.text == 'got') + def command_handler(message): + message.text = message.text + message.text + + tb.process_new_updates([update]) + time.sleep(1) + assert update.message.text == 'got' * 2