Add pattern parameter to events.CallbackQuery (#1212)

telethon/events/callbackquery.py

@@ -19,30 +19,49 @@ class CallbackQuery(EventBuilder):
     `chat_instance` which should be used for inline callbacks.
 
     Args:
-        data (`bytes` | `str` | `callable`, optional):
+        data (`bytes`, `str`, `callable`, optional):
             If set, the inline button payload data must match this data.
             A UTF-8 string can also be given, a regex or a callable. For
             instance, to check against ``'data_1'`` and ``'data_2'`` you
             can use ``re.compile(b'data_')``.
+
+        pattern (`bytes`, `str`, `callable`, `Pattern`, optional):
+            If set, only buttons with payload matching this pattern will be handled.
+            You can specify a regex-like string which will be matched
+            against the payload data, a callable function that returns ``True``
+            if a the payload data is acceptable, or a compiled regex pattern.
+
     """
     def __init__(
-            self, chats=None, *, blacklist_chats=False, func=None, data=None):
+            self, chats=None, *, blacklist_chats=False, func=None, data=None, pattern=None):
         super().__init__(chats, blacklist_chats=blacklist_chats, func=func)
 
-        if isinstance(data, bytes):
+        if data and pattern:
-            self.data = data
+            raise ValueError("Only pass either data or pattern not both.")
-        elif isinstance(data, str):
-            self.data = data.encode('utf-8')
-        elif not data or callable(data):
-            self.data = data
-        elif hasattr(data, 'match') and callable(data.match):
-            if not isinstance(getattr(data, 'pattern', b''), bytes):
-                data = re.compile(data.pattern.encode('utf-8'),
-                                  data.flags & (~re.UNICODE))
 
-            self.data = data.match
+        if isinstance(data, str):
+            data = data.encode('utf-8')
+        if isinstance(pattern, str):
+            pattern = data.encode('utf-8')
+
+        match = data if data else pattern
+
+        if isinstance(match, bytes):
+            self.match = data if data else re.compile(pattern).match
+        elif not match or callable(match):
+            self.match = match
+        elif hasattr(match, 'match') and callable(match.match):
+            if not isinstance(getattr(match, 'pattern', b''), bytes):
+                match = re.compile(match.pattern.encode('utf-8'),
+                                   match.flags & (~re.UNICODE))
+
+            self.match = match.match
         else:
-            raise TypeError('Invalid data type given')
+            raise TypeError('Invalid data or pattern type given')
+
+        self._no_check = all(x is None for x in (
+            self.chats, self.func, self.match,
+        ))
 
     @classmethod
     def build(cls, update, others=None):
@@ -57,21 +76,24 @@ class CallbackQuery(EventBuilder):
 
     def filter(self, event):
         # We can't call super().filter(...) because it ignores chat_instance
+        if self._no_check:
+            return event
+
         if self.chats is not None:
             inside = event.query.chat_instance in self.chats
             if event.chat_id:
                 inside |= event.chat_id in self.chats
 
             if inside == self.blacklist_chats:
-                return None
+                return
 
-        if self.data:
+        if self.match:
-            if callable(self.data):
+            if callable(self.match):
-                event.data_match = self.data(event.query.data)
+                event.data_match = event.pattern_match = self.match(event.query.data)
                 if not event.data_match:
-                    return None
+                    return
-            elif event.query.data != self.data:
+            elif event.query.data != self.match:
-                return None
+                return
 
         if not self.func or self.func(event):
             return event
@@ -88,12 +110,16 @@ class CallbackQuery(EventBuilder):
                 The object returned by the ``data=`` parameter
                 when creating the event builder, if any. Similar
                 to ``pattern_match`` for the new message event.
+            
+            pattern_match (`obj`, optional):
+                Alias for ``data_match``.
         """
         def __init__(self, query, peer, msg_id):
             super().__init__(peer, msg_id=msg_id)
             SenderGetter.__init__(self, query.user_id)
             self.query = query
             self.data_match = None
+            self.pattern_match = None
             self._message = None
             self._answered = False