Merge pull request #785 from bedilbek/middleware-support

Add Middleware Handler

examples/middleware/i18n.py

@@ -0,0 +1,53 @@
+#!/usr/bin/python
+
+# This example shows how to implement i18n (internationalization) l10n (localization) to create
+# multi-language bots with middleware handler.
+#
+# Note: For the sake of simplicity of this example no extra library is used. However, it is recommended to use
+# better i18n systems (gettext and etc) for handling multilingual translations.
+# This is not a working, production-ready sample and it is highly recommended not to use it in production.
+#
+# In this example let's imagine we want to introduce localization or internationalization into our project and
+# we need some global function to activate the language once and to use that language in all other message
+# handler functions for not repeatedly activating it.
+# The middleware (i18n and l10n) is explained:
+
+import telebot
+from telebot import apihelper
+
+apihelper.ENABLE_MIDDLEWARE = True
+
+TRANSLATIONS = {
+    'hello': {
+        'en': 'hello',
+        'ru': 'привет',
+        'uz': 'salom'
+    }
+}
+
+_lang = 'en'
+
+
+def activate(lang):
+    global _lang
+    _lang = lang
+
+
+def _(string):
+    return TRANSLATIONS[string][_lang]
+
+
+bot = telebot.TeleBot('TOKEN')
+
+
+@bot.middleware_handler(update_types=['message'])
+def activate_language(bot_instance, message):
+    activate(message.from_user.language_code)
+
+
+@bot.message_handler(commands=['start'])
+def start(message):
+    bot.send_message(message.chat.id, _('hello'))
+
+
+bot.polling()

examples/middleware/session.py

@@ -0,0 +1,61 @@
+#!/usr/bin/python
+
+# This example shows how to implement session creation and retrieval based on user id with middleware handler.
+#
+# Note: For the sake of simplicity of this example no extra library is used. However, it is recommended to use
+# in-memory or on-disk storage implementations (redis, mysql, postgres and etc) for storing and retrieving structures.
+# This is not a working, production-ready sample and it is highly recommended not to use it in production.
+#
+# In this example let's imagine we want to create a session for each user who communicates with the bot to store
+# different kind of temporary data while session is active. As an example we want to track the state of the user
+# with the help of this session. So, we need a way to store this session data somewhere globally to enable other
+# message handler functions to be able to use it.
+# The middleware session is explained:
+
+import telebot
+from telebot import apihelper
+
+apihelper.ENABLE_MIDDLEWARE = True
+
+INFO_STATE = 'ON_INFO_MENU'
+MAIN_STATE = 'ON_MAIN_MENU'
+
+SESSIONS = {
+    -10000: {
+        'state': INFO_STATE
+    },
+    -11111: {
+        'state': MAIN_STATE
+    }
+}
+
+
+def get_or_create_session(user_id):
+    try:
+        return SESSIONS[user_id]
+    except KeyError:
+        SESSIONS[user_id] = {'state': MAIN_STATE}
+        return SESSIONS[user_id]
+
+
+bot = telebot.TeleBot('TOKEN')
+
+
+@bot.middleware_handler(update_types=['message'])
+def set_session(bot_instance, message):
+    bot_instance.session = get_or_create_session(message.from_user.id)
+
+
+@bot.message_handler(commands=['start'])
+def start(message):
+    bot.session['state'] = MAIN_STATE
+    bot.send_message(message.chat.id, bot.session['state'])
+
+
+@bot.message_handler(commands=['info'])
+def start(message):
+    bot.session['state'] = INFO_STATE
+    bot.send_message(message.chat.id, bot.session['state'])
+
+
+bot.polling()

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,10 @@ class TeleBot:
         new_polls = []
 
         for update in updates:
+
+            if apihelper.ENABLE_MIDDLEWARE:
+                self.process_middlewares(update)
+
             if update.update_id > self.last_update_id:
                 self.last_update_id = update.update_id
             if update.message:
@@ -371,6 +390,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 +1526,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.

telebot/apihelper.py

@@ -26,6 +26,8 @@ FILE_URL = None
 CONNECT_TIMEOUT = 3.5
 READ_TIMEOUT = 9999
 
+ENABLE_MIDDLEWARE = False
+
 
 def _get_req_session(reset=False):
     return util.per_thread('req_session', lambda: requests.session(), reset)

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,43 @@ 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):
+        from telebot import apihelper
+
+        apihelper.ENABLE_MIDDLEWARE = True
+
+        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):
+        from telebot import apihelper
+
+        apihelper.ENABLE_MIDDLEWARE = True
+
+        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