regexp-based inlines and callbacks (szastupov#42)

aiotg/bot.py

@@ -57,9 +57,11 @@ def no_handle(mt):
 
         self._handlers = {mt: no_handle(mt) for mt in MESSAGE_TYPES}
         self._commands = []
-        self._default = lambda c, m: None
-        self._inline = lambda iq: None
-        self._callback = lambda c, cq: None
+        self._callbacks = []
+        self._inlines = []
+        self._default = lambda chat, message: None
+        self._default_callback = lambda chat, cq: None
+        self._default_inline = lambda iq: None
 
     async def loop(self):
         """
@@ -185,6 +187,12 @@ def default(self, callback):
         self._default = callback
         return callback
 
+    def add_inline(self, regexp, fn):
+        """
+        Manually register regexp based callback
+        """
+        self._inlines.append((regexp, fn))
+
     def inline(self, callback):
         """
         Set callback for inline queries
@@ -196,9 +204,29 @@ def inline(self, callback):
         >>>     return iq.answer([
         >>>         {"type": "text", "title": "test", "id", "0"}
         >>>     ])
+
+        >>> @bot.inline(r"myinline-(.+)")
+        >>> def echo(chat, iq, match):
+        >>>     return iq.answer([
+        >>>         {"type": "text", "title": "test", "id", "0"}
+        >>>     ])
         """
-        self._inline = callback
-        return callback
+        if callable(callback):
+            self._default_inline = callback
+            return callback
+        elif isinstance(callback, str):
+            def decorator(fn):
+                self.add_inline(callback, fn)
+                return fn
+            return decorator
+        else:
+            raise TypeError('str expected {} given'.format(type(callback)))
+
+    def add_callback(self, regexp, fn):
+        """
+        Manually register regexp based callback
+        """
+        self._callbacks.append((regexp, fn))
 
     def callback(self, callback):
         """
@@ -209,9 +237,21 @@ def callback(self, callback):
         >>> @bot.callback
         >>> def echo(chat, cq):
         >>>     return cq.answer()
+
+        >>> @bot.callback(r"buttonclick-(.+)")
+        >>> def echo(chat, cq, match):
+        >>>     return chat.reply(match.group(1))
         """
-        self._callback = callback
-        return callback
+        if callable(callback):
+            self._default_callback = callback
+            return callback
+        elif isinstance(callback, str):
+            def decorator(fn):
+                self.add_callback(callback, fn)
+                return fn
+            return decorator
+        else:
+            raise TypeError('str expected {} given'.format(type(callback)))
 
     def handle(self, msg_type):
         """
@@ -486,12 +526,23 @@ def _process_message(self, message):
 
     def _process_inline_query(self, query):
         iq = InlineQuery(self, query)
-        return self._inline(iq)
+
+        for patterns, handler in self._inlines:
+            match = re.search(patterns, query['query'], re.I)
+            if match:
+                return handler(iq, match)
+        return self._default_inline(iq)
 
     def _process_callback_query(self, query):
         chat = Chat.from_message(self, query["message"])
         cq = CallbackQuery(self, query)
-        return self._callback(chat, cq)
+        for patterns, handler in self._callbacks:
+            match = re.search(patterns, cq.data, re.I)
+            if match:
+                return handler(chat, cq, match)
+
+        if not chat.is_group():
+            return self._default_callback(chat, cq)
 
     def _process_updates(self, updates):
         if not updates["ok"]:

examples/callbacks.py

@@ -0,0 +1,39 @@
+
+import json
+from aiotg.ext import ExtendedBot
+
+bot = ExtendedBot(api_token='377308655:AAEUWJSr_0K67SUEQEE57ua5_DEf-TPHZKI')
+
+
+@bot.command(r'/start')
+def start(chat, match):
+
+    markup = {
+            'type': 'InlineKeyboardMarkup',
+            'inline_keyboard': [
+                [{'type': 'InlineKeyboardButton',
+                  'text': 'Button A',
+                  'callback_data': 'buttonclick-A'},
+                 {'type': 'InlineKeyboardButton',
+                  'text': 'Button B',
+                  'callback_data': 'buttonclick-B'}],
+                [{'type': 'InlineKeyboardButton',
+                  'text': 'Nohandle Button',
+                  'callback_data': 'no_callback_data'}],
+                ]
+            }
+
+    chat.send_text('Hello', reply_markup=json.dumps(markup))
+
+
+@bot.callback(r'buttonclick-(\w+)')
+def buttonclick(chat, cq, match):
+    chat.send_text('You clicked {}'.format(match.group(1)))
+
+
+@bot.default_callback
+def unhandled_callbacks(chat, cq):
+    chat.send_text('Unhandled callback fired')
+
+
+bot.run(debug=True)

tests/test_callbacks.py

@@ -6,6 +6,7 @@
 from aiotg.mock import MockBot
 from testfixtures import LogCapture
 
+
 API_TOKEN = "test_token"
 bot = Bot(API_TOKEN)
 
@@ -68,7 +69,7 @@ def default(chat, message):
     assert called_with == "foo bar"
 
 
-def test_inline():
+def test_default_inline():
     called_with = None
 
     @bot.inline
@@ -80,11 +81,23 @@ def inline(query):
     assert called_with == "foo bar"
 
 
+def test_inline():
+    called_with = None
+
+    @bot.inline(r'query-(\w+)')
+    def inline(query, match):
+        nonlocal called_with
+        called_with = match.group(1)
+
+    bot._process_inline_query(inline_query("query-foo"))
+    assert called_with == "foo"
+
+
 def test_callback_default():
     bot._process_callback_query(callback_query("foo"))
 
 
-def test_callback():
+def test_default_callback():
     called_with = None
 
     @bot.callback
@@ -96,6 +109,18 @@ def callback(chat, cq):
     assert called_with == "foo"
 
 
+def test_callback():
+    called_with = None
+
+    @bot.callback(r'click-(\w+)')
+    def click_callback(chat, cq, match):
+        nonlocal called_with
+        called_with = match.group(1)
+
+    bot._process_callback_query(callback_query("click-foo"))
+    assert called_with == "foo"
+
+
 def test_updates():
     update = {
         "update_id": 0,