Allow adding callback methods to UpdateState

telethon/telegram_client.py

@@ -907,16 +907,13 @@ class TelegramClient(TelegramBareClient):
     def add_update_handler(self, handler):
         """Adds an update handler (a function which takes a TLObject,
           an update, as its parameter) and listens for updates"""
-        return  # TODO Implement
-        self._update_callbacks.append(handler)
+        self.updates.handlers.append(handler)
 
     def remove_update_handler(self, handler):
-        return  # TODO Implement
-        self._update_callbacks.remove(handler)
+        self.updates.handlers.remove(handler)
 
     def list_update_handlers(self):
-        return  # TODO Implement
-        return self._update_callbacks[:]
+        return self.updates.handlers[:]
 
     # endregion
 

telethon/update_state.py

@@ -8,6 +8,7 @@ class UpdateState:
     """
     def __init__(self, enabled):
         self.enabled = enabled
+        self.handlers = []
         self._updates_lock = Lock()
         self._updates_available = Event()
         self._updates = deque()
@@ -30,9 +31,11 @@ class UpdateState:
         """Processes an update object. This method is normally called by
            the library itself.
         """
-        if not self.enabled:
-            return
+        for handler in self.handlers:
+            handler(update)
+
+        if self.enabled:
+            with self._updates_lock:
+                self._updates.append(update)
+                self._updates_available.set()
 
-        with self._updates_lock:
-            self._updates.append(update)
-            self._updates_available.set()