Merge branch 'remove-handler-and-locks' of https://github.com/JosXa/pyrogram into JosXa-remove-handler-and-locks

# Conflicts:
#	pyrogram/client/dispatcher/dispatcher.py

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.

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:
+        with self._handler_lock:
-            self.groups[group] = []
+            if group not in self.groups:
-            self.groups = OrderedDict(sorted(self.groups.items()))
+                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():
+        with self._handler_lock:
-            for handler in group:
+            for group in self.groups.values():
-                if is_raw:
+                for handler in group:
-                    if not isinstance(handler, RawUpdateHandler):
+                    if is_raw:
-                        continue
+                        if not isinstance(handler, RawUpdateHandler):
-
-                    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):
                             continue
 
-                        args = (self.client, message)
+                        args = (self.client, update, users, chats)
-                    elif callback_query and isinstance(handler, CallbackQueryHandler):
-                        if not handler.check(callback_query):
-                            continue
-
-                        args = (self.client, callback_query)
                     else:
-                        continue
+                        message = (update.message
+                                   or update.channel_post
+                                   or update.edited_message
+                                   or update.edited_channel_post)
 
-                handler.callback(*args)
+                        callback_query = update.callback_query
-                break
+
+                        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