From d23737520880010e6de068e6f600f2023e941ebd Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 7 Sep 2017 18:58:54 +0200
Subject: [PATCH] Allow adding callback methods to UpdateState

---
 telethon/telegram_client.py |  9 +++------
 telethon/update_state.py    | 13 ++++++++-----
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py
index ef83cdaa..545633e7 100644
--- a/telethon/telegram_client.py
+++ b/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
 
diff --git a/telethon/update_state.py b/telethon/update_state.py
index 14d362a9..fd1a0935 100644
--- a/telethon/update_state.py
+++ b/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()