refactor: optimize update dispatching and filter performance

Author: SychO3

pyrogram/dispatcher.py

@@ -262,11 +262,25 @@ async def deleted_business_messages_parser(update, users, chats):
     async def start(self):
         if callable(self.client.start_handler):
             try:
-                await self.client.start_handler(self.client)
+                if inspect.iscoroutinefunction(self.client.start_handler):
+                    await self.client.start_handler(self.client)
+                else:
+                    result = self.client.start_handler(self.client)
+                    if inspect.isawaitable(result):
+                        await result
             except Exception as e:
-                log.exception(e)
+                log.exception("start_handler raised: %s", e)
 
         if not self.client.no_updates:
+            self.locks_list.clear()
+            self.handler_worker_tasks.clear()
+
+            if 0 not in self.groups or self.conversation_handler not in self.groups.get(0, []):
+                if 0 not in self.groups:
+                    self.groups[0] = []
+                    self.groups = OrderedDict(sorted(self.groups.items()))
+                self.groups[0].insert(0, self.conversation_handler)
+
             for i in range(self.client.workers):
                 self.locks_list.append(asyncio.Lock())
 
@@ -301,7 +315,7 @@ async def stop(self, clear_handlers: bool = True):
                         self.client.remove_listener(listener)
                         if getattr(listener, "future", None) and not listener.future.done():
                             try:
-                                listener.future.set_exception(asyncio.CancelledError())
+                                listener.future.cancel()
                             except Exception:
                                 pass
                     except Exception:
@@ -311,9 +325,14 @@ async def stop(self, clear_handlers: bool = True):
 
         if callable(self.client.stop_handler):
             try:
-                await self.client.stop_handler(self.client)
+                if inspect.iscoroutinefunction(self.client.stop_handler):
+                    await self.client.stop_handler(self.client)
+                else:
+                    result = self.client.stop_handler(self.client)
+                    if inspect.isawaitable(result):
+                        await result
             except Exception as e:
-                log.exception(e)
+                log.exception("stop_handler raised: %s", e)
 
         if not self.client.no_updates:
             for i in range(self.client.workers):
@@ -322,8 +341,10 @@ async def stop(self, clear_handlers: bool = True):
             for i in self.handler_worker_tasks:
                 await i
 
+            self.handler_worker_tasks.clear()
+            self.locks_list.clear()
+
             if clear_handlers:
-                self.handler_worker_tasks.clear()
                 self.groups.clear()
 
             log.info("Stopped %s HandlerTasks", self.client.workers)
@@ -339,6 +360,8 @@ async def fn():
                     self.groups = OrderedDict(sorted(self.groups.items()))
 
                 self.groups[group].append(handler)
+            except Exception as e:
+                log.exception("Failed to add handler: %s", e)
             finally:
                 for lock in self.locks_list:
                     lock.release()
@@ -352,14 +375,17 @@ async def fn():
 
             try:
                 if group not in self.groups:
-                    raise ValueError(
-                        f"Group {group} does not exist. Handler was not removed."
-                    )
+                    log.warning("Group %s does not exist. Handler was not removed.", group)
+                    return
 
                 self.groups[group].remove(handler)
 
                 if not self.groups[group]:
                     del self.groups[group]
+            except ValueError:
+                log.warning("Handler not found in group %s.", group)
+            except Exception as e:
+                log.exception("Failed to remove handler: %s", e)
             finally:
                 for lock in self.locks_list:
                     lock.release()
@@ -395,16 +421,16 @@ async def handler_worker(self, lock):
                                 try:
                                     if await handler.check(self.client, parsed_update):
                                         args = (parsed_update,)
-                                except Exception as e:
-                                    log.exception(e)
+                                except Exception:
+                                    log.exception("Handler check failed")
                                     continue
 
                             elif isinstance(handler, RawUpdateHandler):
                                 try:
                                     if await handler.check(self.client, update):
                                         args = (update, users, chats)
-                                except Exception as e:
-                                    log.exception(e)
+                                except Exception:
+                                    log.exception("Raw handler check failed")
                                     continue
 
                             if args is None:
@@ -421,8 +447,7 @@ async def handler_worker(self, lock):
                                         *args
                                     )
                             except asyncio.CancelledError:
-                                # Swallow task cancellations during shutdown/interrupt
-                                pass
+                                raise
                             except pyrogram.StopPropagation:
                                 raise
                             except pyrogram.ContinuePropagation:
@@ -435,8 +460,8 @@ async def handler_worker(self, lock):
                             break
             except pyrogram.StopPropagation:
                 pass
-            except Exception as e:
-                log.exception(e)
+            except Exception:
+                log.exception("Unhandled exception in handler worker")
 
     async def handle_update_handler_exception(
         self,

pyrogram/filters.py

@@ -43,9 +43,10 @@ def __or__(self, other):
 class InvertFilter(Filter):
     def __init__(self, base):
         self.base = base
+        self._base_is_async = inspect.iscoroutinefunction(base.__call__)
 
     async def __call__(self, client: "pyrogram.Client", update: Update):
-        if inspect.iscoroutinefunction(self.base.__call__):
+        if self._base_is_async:
             x = await self.base(client, update)
         else:
             x = await client.loop.run_in_executor(
@@ -61,9 +62,11 @@ class AndFilter(Filter):
     def __init__(self, base, other):
         self.base = base
         self.other = other
+        self._base_is_async = inspect.iscoroutinefunction(base.__call__)
+        self._other_is_async = inspect.iscoroutinefunction(other.__call__)
 
     async def __call__(self, client: "pyrogram.Client", update: Update):
-        if inspect.iscoroutinefunction(self.base.__call__):
+        if self._base_is_async:
             x = await self.base(client, update)
         else:
             x = await client.loop.run_in_executor(
@@ -72,11 +75,10 @@ async def __call__(self, client: "pyrogram.Client", update: Update):
                 client, update
             )
 
-        # short circuit
         if not x:
             return False
 
-        if inspect.iscoroutinefunction(self.other.__call__):
+        if self._other_is_async:
             y = await self.other(client, update)
         else:
             y = await client.loop.run_in_executor(
@@ -85,16 +87,18 @@ async def __call__(self, client: "pyrogram.Client", update: Update):
                 client, update
             )
 
-        return x and y
+        return y
 
 
 class OrFilter(Filter):
     def __init__(self, base, other):
         self.base = base
         self.other = other
+        self._base_is_async = inspect.iscoroutinefunction(base.__call__)
+        self._other_is_async = inspect.iscoroutinefunction(other.__call__)
 
     async def __call__(self, client: "pyrogram.Client", update: Update):
-        if inspect.iscoroutinefunction(self.base.__call__):
+        if self._base_is_async:
             x = await self.base(client, update)
         else:
             x = await client.loop.run_in_executor(
@@ -103,11 +107,10 @@ async def __call__(self, client: "pyrogram.Client", update: Update):
                 client, update
             )
 
-        # short circuit
         if x:
             return True
 
-        if inspect.iscoroutinefunction(self.other.__call__):
+        if self._other_is_async:
             y = await self.other(client, update)
         else:
             y = await client.loop.run_in_executor(
@@ -116,7 +119,7 @@ async def __call__(self, client: "pyrogram.Client", update: Update):
                 client, update
             )
 
-        return x or y
+        return y
 
 
 CUSTOM_FILTER_NAME = "CustomFilter"
@@ -146,7 +149,7 @@ def create(func: Callable, name: str = None, **kwargs) -> Filter:
             :meth:`~pyrogram.filters.command` or :meth:`~pyrogram.filters.regex`.
     """
     return type(
-        name or func.__name__ or CUSTOM_FILTER_NAME,
+        name or getattr(func, "__name__", None) or CUSTOM_FILTER_NAME,
         (Filter,),
         {"__call__": func, **kwargs}
     )()
@@ -923,11 +926,11 @@ async def paid_message_filter(_, __, m: Message):
 
 # region linked_channel_filter
 async def linked_channel_filter(_, __, m: Message):
-    return bool(
-        m.forward_origin and
-        m.forward_origin.type == enums.MessageOriginType.CHANNEL and
-        m.forward_origin.chat == m.sender_chat
-    )
+    origin = m.forward_origin
+    if not (origin and origin.type == enums.MessageOriginType.CHANNEL and m.sender_chat):
+        return False
+    origin_chat = getattr(origin, "chat", None)
+    return origin_chat is not None and origin_chat.id == m.sender_chat.id
 
 
 linked_channel = create(linked_channel_filter)
@@ -999,7 +1002,7 @@ def command(commands: Union[str, List[str]], prefixes: Optional[Union[str, List[
     command_re = re.compile(r"([\"'])(.*?)(?<!\\)\1|(\S+)")
 
     async def func(flt, client: pyrogram.Client, message: Message):
-        username = client.me.username or ""
+        username = client.me.username if client.me else ""
         text = message.text or message.caption
         message.command = None
 
@@ -1013,11 +1016,12 @@ async def func(flt, client: pyrogram.Client, message: Message):
             without_prefix = text[len(prefix):]
 
             for cmd in flt.commands:
-                if not re.match(rf"^(?:{cmd}(?:@?{username})?)(?:\s|$)", without_prefix,
+                escaped_cmd = re.escape(cmd)
+                if not re.match(rf"^(?:{escaped_cmd}(?:@?{username})?)(?:\s|$)", without_prefix,
                                 flags=re.IGNORECASE if not flt.case_sensitive else 0):
                     continue
 
-                without_command = re.sub(rf"{cmd}(?:@?{username})?\s?", "", without_prefix, count=1,
+                without_command = re.sub(rf"{escaped_cmd}(?:@?{username})?\s?", "", without_prefix, count=1,
                                          flags=re.IGNORECASE if not flt.case_sensitive else 0)
 
                 # match.groups are 1-indexed, group(1) is the quote, group(2) is the text
@@ -1087,6 +1091,8 @@ async def func(flt, _, update: Update):
 
         if value:
             update.matches = list(flt.p.finditer(value)) or None
+        else:
+            update.matches = None
 
         return bool(update.matches)
 

pyrogram/handlers/callback_query_handler.py

@@ -112,7 +112,7 @@ async def check_if_has_matching_listener(
                     listener_does_match = await filters(client, query)
                 else:
                     listener_does_match = await client.loop.run_in_executor(
-                        None, filters, client, query
+                        client.executor, filters, client, query
                     )
             else:
                 listener_does_match = True
@@ -132,32 +132,30 @@ async def check(self, client: "pyrogram.Client", query: CallbackQuery):
             client, query
         )
 
+        query._matched_listener = listener if listener_does_match else None
+
         if callable(self.filters):
             if iscoroutinefunction(self.filters.__call__):
                 handler_does_match = await self.filters(client, query)
             else:
                 handler_does_match = await client.loop.run_in_executor(
-                    None, self.filters, client, query
+                    client.executor, self.filters, client, query
                 )
         else:
             handler_does_match = True
 
-        data = self.compose_data_identifier(query)
-
-        if PyromodConfig.unallowed_click_alert:
-            # matches with the current query but from any user
+        if PyromodConfig.unallowed_click_alert and listener:
+            data = self.compose_data_identifier(query)
             permissive_identifier = Identifier(
                 chat_id=data.chat_id,
                 message_id=data.message_id,
                 inline_message_id=data.inline_message_id,
                 from_user_id=None,
             )
 
-            matches = permissive_identifier.matches(data)
-
             if (
-                listener
-                and (matches and not listener_does_match)
+                permissive_identifier.matches(data)
+                and not listener_does_match
                 and listener.unallowed_click_alert
             ):
                 alert = (
@@ -168,8 +166,6 @@ async def check(self, client: "pyrogram.Client", query: CallbackQuery):
                 await query.answer(alert)
                 return False
 
-        # let handler get the chance to handle if listener
-        # exists but its filters doesn't match
         return listener_does_match or handler_does_match
 
     async def resolve_future_or_callback(
@@ -183,9 +179,11 @@ async def resolve_future_or_callback(
         :param args: The arguments to call the callback with.
         :return: None
         """
-        listener_does_match, listener = await self.check_if_has_matching_listener(
-            client, query
-        )
+        listener = getattr(query, '_matched_listener', None)
+        listener_does_match = listener is not None
+
+        if not listener_does_match:
+            listener_does_match, listener = await self.check_if_has_matching_listener(client, query)
 
         if listener and listener_does_match:
             client.remove_listener(listener)

pyrogram/handlers/conversation_handler.py

@@ -29,11 +29,20 @@
 class ConversationHandler(MessageHandler, CallbackQueryHandler):
     """The Conversation handler class."""
     def __init__(self):
-        # Initialize base handler with a no-op async callback to satisfy handler interface
         Handler.__init__(self, self.callback)
+        self.original_callback = self.callback
         self.waiters = {}
 
+    def register_waiter(self, chat_id, waiter):
+        old = self.waiters.get(chat_id)
+        if old and not old['future'].done():
+            old['future'].cancel()
+        self.waiters[chat_id] = waiter
+
     async def check(self, client: "pyrogram.Client", update: Union[Message, CallbackQuery]):
+        if not self.waiters:
+            return False
+
         if isinstance(update, Message) and update.outgoing:
             return False
 
@@ -71,8 +80,9 @@ async def check(self, client: "pyrogram.Client", update: Union[Message, Callback
 
     @staticmethod
     async def callback(_, __):
-        pass
+        raise pyrogram.StopPropagation
 
     def delete_waiter(self, chat_id, future):
-        if future == self.waiters[chat_id]['future']:
+        waiter = self.waiters.get(chat_id)
+        if waiter and waiter.get('future') == future:
             del self.waiters[chat_id]