diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 35124c0f..22eee1bf 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -262,12 +262,13 @@ class Client: return func return decorator - + + # TODO: Maybe make add_handler return (handler, group)? def add_handler(self, handler, group: int = 0): - """Use this method to register an event handler. + """Use this method to register an update handler. You can register multiple handlers, but at most one handler within a group - will be used for a single event. To handle the same event more than once, register + will be used for a single update. To handle the same update more than once, register your handler using a different group id (lower group id == higher priority). Args: @@ -279,6 +280,22 @@ class Client: """ self.dispatcher.add_handler(handler, group) + def remove_handler(self, handler, group: int = 0): + """Removes a previously-added update handler. + + Make sure to provide the right group that the handler was added in. You can use + the return value of the :meth:`add_handler` method, a tuple of (handler, group), and + pass it directly. + + Args: + handler (``Handler``): + The handler to be removed. + + group (``int``, optional): + The group identifier, defaults to 0. + """ + self.dispatcher.remove_handler(handler, group) + def start(self, debug: bool = False): """Use this method to start the Client after creating it. Requires no parameters. diff --git a/pyrogram/client/dispatcher/dispatcher.py b/pyrogram/client/dispatcher/dispatcher.py index 15e6ccd1..a9439d6e 100644 --- a/pyrogram/client/dispatcher/dispatcher.py +++ b/pyrogram/client/dispatcher/dispatcher.py @@ -50,6 +50,8 @@ class Dispatcher: self.updates = Queue() self.groups = OrderedDict() + self._handler_lock = threading.Lock() + def start(self): for i in range(self.workers): self.workers_list.append( @@ -71,43 +73,52 @@ class Dispatcher: self.workers_list.clear() def add_handler(self, handler, group: int): - if group not in self.groups: - self.groups[group] = [] - self.groups = OrderedDict(sorted(self.groups.items())) + with self._handler_lock: + if group not in self.groups: + self.groups[group] = [] + self.groups = OrderedDict(sorted(self.groups.items())) - self.groups[group].append(handler) + self.groups[group].append(handler) + + def remove_handler(self, handler, group: int): + with self._handler_lock: + if group not in self.groups: + raise ValueError("Group {} does not exist. " + "Handler was not removed.".format(group)) + self.groups[group].remove(handler) def dispatch(self, update, users: dict = None, chats: dict = None, is_raw: bool = False): - for group in self.groups.values(): - for handler in group: - if is_raw: - if not isinstance(handler, RawUpdateHandler): - continue - - args = (self.client, update, users, chats) - else: - message = (update.message - or update.channel_post - or update.edited_message - or update.edited_channel_post) - - callback_query = update.callback_query - - if message and isinstance(handler, MessageHandler): - if not handler.check(message): + with self._handler_lock: + for group in self.groups.values(): + for handler in group: + if is_raw: + if not isinstance(handler, RawUpdateHandler): continue - args = (self.client, message) - elif callback_query and isinstance(handler, CallbackQueryHandler): - if not handler.check(callback_query): - continue - - args = (self.client, callback_query) + args = (self.client, update, users, chats) else: - continue + message = (update.message + or update.channel_post + or update.edited_message + or update.edited_channel_post) - handler.callback(*args) - break + callback_query = update.callback_query + + if message and isinstance(handler, MessageHandler): + if not handler.check(message): + continue + + args = (self.client, message) + elif callback_query and isinstance(handler, CallbackQueryHandler): + if not handler.check(callback_query): + continue + + args = (self.client, callback_query) + else: + continue + + handler.callback(*args) + break def update_worker(self): name = threading.current_thread().name