From 36dabc4928543f0ecf2a30021ba4d7aad8c3e526 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Sun, 1 Oct 2017 16:02:29 +0200
Subject: [PATCH] Process entities from the MtProtoSender instead
 TelegramBareClient

This way, update objects will also be processed when they occur.
---
 telethon/network/mtproto_sender.py | 11 +++--------
 telethon/telegram_bare_client.py   | 19 ++++---------------
 telethon/tl/session.py             | 21 +++++++++++++++++++--
 3 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/telethon/network/mtproto_sender.py b/telethon/network/mtproto_sender.py
index 6558a20c..147f4bef 100644
--- a/telethon/network/mtproto_sender.py
+++ b/telethon/network/mtproto_sender.py
@@ -213,14 +213,8 @@ class MtProtoSender:
         # If the code is not parsed manually then it should be a TLObject.
         if code in tlobjects:
             result = reader.tgread_object()
-            if state is None:
-                self._logger.debug(
-                    'Ignoring unhandled TLObject %s', repr(result)
-                )
-            else:
-                self._logger.debug(
-                    'Processing TLObject %s', repr(result)
-                )
+            self.session.process_entities(result)
+            if state:
                 state.process(result)
 
             return True
@@ -364,6 +358,7 @@ class MtProtoSender:
                     reader.seek(-4)
                     request.on_response(reader)
 
+                self.session.process_entities(request.result)
                 request.confirm_received.set()
                 return True
             else:
diff --git a/telethon/telegram_bare_client.py b/telethon/telegram_bare_client.py
index 6174ecc1..60e6d866 100644
--- a/telethon/telegram_bare_client.py
+++ b/telethon/telegram_bare_client.py
@@ -494,21 +494,10 @@ class TelegramBareClient:
                 # rejected by the other party as a whole."
                 return None
 
-            # Save all input entities we know of
-            entities = []
-            results = []
-            for x in requests:
-                y = x.result
-                results.append(y)
-                if hasattr(y, 'chats') and hasattr(y.chats, '__iter__'):
-                    entities.extend(y.chats)
-                if hasattr(y, 'users') and hasattr(y.users, '__iter__'):
-                    entities.extend(y.users)
-
-            if self.session.add_entities(entities):
-                self.session.save()  # Save if any new entities got added
-
-            return results[0] if len(results) == 1 else results
+            if len(requests) == 1:
+                return requests[0].result
+            else:
+                return [x.result for x in requests]
 
         except (PhoneMigrateError, NetworkMigrateError,
                 UserMigrateError) as e:
diff --git a/telethon/tl/session.py b/telethon/tl/session.py
index 41a6a61c..d3854c8d 100644
--- a/telethon/tl/session.py
+++ b/telethon/tl/session.py
@@ -185,11 +185,28 @@ class Session:
         correct = correct_msg_id >> 32
         self.time_offset = correct - now
 
+    def process_entities(self, tlobject):
+        """Processes all the found entities on the given TLObject,
+           unless .save_entities is False, and saves the session file.
+        """
+        if not self.save_entities:
+            return
+
+        # Save all input entities we know of
+        entities = []
+        if hasattr(tlobject, 'chats') and hasattr(tlobject.chats, '__iter__'):
+            entities.extend(tlobject.chats)
+        if hasattr(tlobject, 'users') and hasattr(tlobject.users, '__iter__'):
+            entities.extend(tlobject.users)
+
+        if self.add_entities(entities):
+            self.save()  # Save if any new entities got added
+
     def add_entities(self, entities):
-        """Adds new input entities to the local database of them.
+        """Adds new input entities to the local database unconditionally.
            Unknown types will be ignored.
         """
-        if not entities or not self.save_entities:
+        if not entities:
             return False
 
         new = {}