From 5be9df0eece36d4f5ea29447edec309991f443c1 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Wed, 4 Oct 2017 13:58:36 +0200
Subject: [PATCH 01/11] Add a basic EntityDatabase class

---
 telethon/entity_database.py | 75 +++++++++++++++++++++++++++++++++++++
 1 file changed, 75 insertions(+)
 create mode 100644 telethon/entity_database.py

diff --git a/telethon/entity_database.py b/telethon/entity_database.py
new file mode 100644
index 00000000..1e96fe51
--- /dev/null
+++ b/telethon/entity_database.py
@@ -0,0 +1,75 @@
+from . import utils
+from .tl import TLObject
+
+
+class EntityDatabase:
+    def __init__(self, enabled=True):
+        self.enabled = enabled
+
+        self._entities = {}  # marked_id: user|chat|channel
+
+        # TODO Allow disabling some extra mappings
+        self._username_id = {}  # username: marked_id
+
+    def add(self, entity):
+        if not self.enabled:
+            return
+
+        # Adds or updates the given entity
+        marked_id = utils.get_peer_id(entity, add_mark=True)
+        try:
+            old_entity = self._entities[marked_id]
+            old_entity.__dict__.update(entity)  # Keep old references
+
+            # Update must delete old username
+            username = getattr(old_entity, 'username', None)
+            if username:
+                del self._username_id[username.lower()]
+        except KeyError:
+            # Add new entity
+            self._entities[marked_id] = entity
+
+        # Always update username if any
+        username = getattr(entity, 'username', None)
+        if username:
+            self._username_id[username.lower()] = marked_id
+
+    def __getitem__(self, key):
+        """Accepts a digit only string as phone number,
+           otherwise it's treated as an username.
+
+           If an integer is given, it's treated as the ID of the desired User.
+           The ID given won't try to be guessed as the ID of a chat or channel,
+           as there may be an user with that ID, and it would be unreliable.
+
+           If a Peer is given (PeerUser, PeerChat, PeerChannel),
+           its specific entity is retrieved as User, Chat or Channel.
+           Note that megagroups are channels with .megagroup = True.
+        """
+        if isinstance(key, str):
+            # TODO Parse phone properly, currently only usernames
+            key = key.lstrip('@').lower()
+            # TODO Use the client to return from username if not found
+            return self._entities[self._username_id[key]]
+
+        if isinstance(key, int):
+            return self._entities[key]  # normal IDs are assumed users
+
+        if isinstance(key, TLObject) and type(key).SUBCLASS_OF_ID == 0x2d45687:
+            return self._entities[utils.get_peer_id(key, add_mark=True)]
+
+        raise KeyError(key)
+
+    def __delitem__(self, key):
+        target = self[key]
+        del self._entities[key]
+        if getattr(target, 'username'):
+            del self._username_id[target.username]
+
+    # TODO Allow search by name by tokenizing the input and return a list
+
+    def clear(self, target=None):
+        if target is None:
+            self._entities.clear()
+        else:
+            del self[target]

From a0fc5ed54e813d48eee00424b171ba4d6885a17e Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Wed, 4 Oct 2017 21:02:45 +0200
Subject: [PATCH 02/11] Use EntityDatabase in the Session class

---
 telethon/entity_database.py    |  75 ------------------
 telethon/tl/entity_database.py | 140 +++++++++++++++++++++++++++++++++
 telethon/tl/session.py         |  70 ++---------------
 3 files changed, 146 insertions(+), 139 deletions(-)
 delete mode 100644 telethon/entity_database.py
 create mode 100644 telethon/tl/entity_database.py

diff --git a/telethon/entity_database.py b/telethon/entity_database.py
deleted file mode 100644
index 1e96fe51..00000000
--- a/telethon/entity_database.py
+++ /dev/null
@@ -1,75 +0,0 @@
-from . import utils
-from .tl import TLObject
-
-
-class EntityDatabase:
-    def __init__(self, enabled=True):
-        self.enabled = enabled
-
-        self._entities = {}  # marked_id: user|chat|channel
-
-        # TODO Allow disabling some extra mappings
-        self._username_id = {}  # username: marked_id
-
-    def add(self, entity):
-        if not self.enabled:
-            return
-
-        # Adds or updates the given entity
-        marked_id = utils.get_peer_id(entity, add_mark=True)
-        try:
-            old_entity = self._entities[marked_id]
-            old_entity.__dict__.update(entity)  # Keep old references
-
-            # Update must delete old username
-            username = getattr(old_entity, 'username', None)
-            if username:
-                del self._username_id[username.lower()]
-        except KeyError:
-            # Add new entity
-            self._entities[marked_id] = entity
-
-        # Always update username if any
-        username = getattr(entity, 'username', None)
-        if username:
-            self._username_id[username.lower()] = marked_id
-
-    def __getitem__(self, key):
-        """Accepts a digit only string as phone number,
-           otherwise it's treated as an username.
-
-           If an integer is given, it's treated as the ID of the desired User.
-           The ID given won't try to be guessed as the ID of a chat or channel,
-           as there may be an user with that ID, and it would be unreliable.
-
-           If a Peer is given (PeerUser, PeerChat, PeerChannel),
-           its specific entity is retrieved as User, Chat or Channel.
-           Note that megagroups are channels with .megagroup = True.
-        """
-        if isinstance(key, str):
-            # TODO Parse phone properly, currently only usernames
-            key = key.lstrip('@').lower()
-            # TODO Use the client to return from username if not found
-            return self._entities[self._username_id[key]]
-
-        if isinstance(key, int):
-            return self._entities[key]  # normal IDs are assumed users
-
-        if isinstance(key, TLObject) and type(key).SUBCLASS_OF_ID == 0x2d45687:
-            return self._entities[utils.get_peer_id(key, add_mark=True)]
-
-        raise KeyError(key)
-
-    def __delitem__(self, key):
-        target = self[key]
-        del self._entities[key]
-        if getattr(target, 'username'):
-            del self._username_id[target.username]
-
-    # TODO Allow search by name by tokenizing the input and return a list
-
-    def clear(self, target=None):
-        if target is None:
-            self._entities.clear()
-        else:
-            del self[target]
diff --git a/telethon/tl/entity_database.py b/telethon/tl/entity_database.py
new file mode 100644
index 00000000..6a1d7dbb
--- /dev/null
+++ b/telethon/tl/entity_database.py
@@ -0,0 +1,140 @@
+from threading import Lock
+
+from .. import utils
+from ..tl import TLObject
+from ..tl.types import User, Chat, Channel
+
+
+class EntityDatabase:
+    def __init__(self, input_list=None, enabled=True):
+        self.enabled = enabled
+
+        self._lock = Lock()
+        self._entities = {}  # marked_id: user|chat|channel
+
+        if input_list:
+            self._input_entities = {k: v for k, v in input_list}
+        else:
+            self._input_entities = {}  # marked_id: hash
+
+        # TODO Allow disabling some extra mappings
+        self._username_id = {}  # username: marked_id
+
+    def process(self, tlobject):
+        """Processes all the found entities on the given TLObject,
+           unless .enabled is False.
+
+           Returns True if new input entities were added.
+        """
+        if not self.enabled:
+            return False
+
+        # 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)
+
+        return self.expand(entities)
+
+    def expand(self, entities):
+        """Adds new input entities to the local database unconditionally.
+           Unknown types will be ignored.
+        """
+        if not entities or not self.enabled:
+            return False
+
+        new = []  # Array of entities (User, Chat, or Channel)
+        new_input = {}  # Dictionary of {entity_marked_id: access_hash}
+        for e in entities:
+            if not isinstance(e, TLObject):
+                continue
+
+            try:
+                p = utils.get_input_peer(e)
+                new_input[utils.get_peer_id(p, add_mark=True)] = \
+                    getattr(p, 'access_hash', 0)  # chats won't have hash
+
+                if isinstance(e, User) \
+                        or isinstance(e, Chat) \
+                        or isinstance(e, Channel):
+                    new.append(e)
+            except ValueError:
+                pass
+
+        with self._lock:
+            before = len(self._input_entities)
+            self._input_entities.update(new_input)
+            for e in new:
+                self._add_full_entity(e)
+            return len(self._input_entities) != before
+
+    def _add_full_entity(self, entity):
+        """Adds a "full" entity (User, Chat or Channel, not "Input*").
+
+           Not to be confused with UserFull, ChatFull, or ChannelFull,
+           "full" means simply not "Input*".
+        """
+        marked_id = utils.get_peer_id(
+            utils.get_input_peer(entity), add_mark=True
+        )
+        try:
+            old_entity = self._entities[marked_id]
+            old_entity.__dict__.update(entity.__dict__)  # Keep old references
+
+            # Update must delete old username
+            username = getattr(old_entity, 'username', None)
+            if username:
+                del self._username_id[username.lower()]
+        except KeyError:
+            # Add new entity
+            self._entities[marked_id] = entity
+
+        # Always update username if any
+        username = getattr(entity, 'username', None)
+        if username:
+            self._username_id[username.lower()] = marked_id
+
+    def __getitem__(self, key):
+        """Accepts a digit only string as phone number,
+           otherwise it's treated as an username.
+
+           If an integer is given, it's treated as the ID of the desired User.
+           The ID given won't try to be guessed as the ID of a chat or channel,
+           as there may be an user with that ID, and it would be unreliable.
+
+           If a Peer is given (PeerUser, PeerChat, PeerChannel),
+           its specific entity is retrieved as User, Chat or Channel.
+           Note that megagroups are channels with .megagroup = True.
+        """
+        if isinstance(key, str):
+            # TODO Parse phone properly, currently only usernames
+            key = key.lstrip('@').lower()
+            # TODO Use the client to return from username if not found
+            return self._entities[self._username_id[key]]
+
+        if isinstance(key, int):
+            return self._entities[key]  # normal IDs are assumed users
+
+        if isinstance(key, TLObject) and type(key).SUBCLASS_OF_ID == 0x2d45687:
+            return self._entities[utils.get_peer_id(key, add_mark=True)]
+
+        raise KeyError(key)
+
+    def __delitem__(self, key):
+        target = self[key]
+        del self._entities[key]
+        if getattr(target, 'username'):
+            del self._username_id[target.username]
+
+    # TODO Allow search by name by tokenizing the input and return a list
+
+    def get_input_list(self):
+        return list(self._input_entities.items())
+
+    def clear(self, target=None):
+        if target is None:
+            self._entities.clear()
+        else:
+            del self[target]
diff --git a/telethon/tl/session.py b/telethon/tl/session.py
index d3854c8d..2b691ad7 100644
--- a/telethon/tl/session.py
+++ b/telethon/tl/session.py
@@ -6,11 +6,8 @@ from base64 import b64encode, b64decode
 from os.path import isfile as file_exists
 from threading import Lock
 
-from .. import helpers, utils
-from ..tl.types import (
-    InputPeerUser, InputPeerChat, InputPeerChannel,
-    PeerUser, PeerChat, PeerChannel
-)
+from .entity_database import EntityDatabase
+from .. import helpers
 
 
 class Session:
@@ -70,8 +67,7 @@ class Session:
         self.auth_key = None
         self.layer = 0
         self.salt = 0  # Unsigned long
-        self._input_entities = {}  # {marked_id: hash}
-        self._entities_lock = Lock()
+        self.entities = EntityDatabase()  # Known and cached entities
 
     def save(self):
         """Saves the current session object as session_user_id.session"""
@@ -90,7 +86,7 @@ class Session:
                         if self.auth_key else None
                 }
                 if self.save_entities:
-                    out_dict['entities'] = list(self._input_entities.items())
+                    out_dict['entities'] = self.entities.get_input_list()
 
                 json.dump(out_dict, file)
 
@@ -139,8 +135,7 @@ class Session:
                         key = b64decode(data['auth_key_data'])
                         result.auth_key = AuthKey(data=key)
 
-                    for e_mid, e_hash in data.get('entities', []):
-                        result._input_entities[e_mid] = e_hash
+                    result.entities = EntityDatabase(data.get('entities', []))
 
             except (json.decoder.JSONDecodeError, UnicodeDecodeError):
                 pass
@@ -186,58 +181,5 @@ class Session:
         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):
+        if self.entities.process(tlobject):
             self.save()  # Save if any new entities got added
-
-    def add_entities(self, entities):
-        """Adds new input entities to the local database unconditionally.
-           Unknown types will be ignored.
-        """
-        if not entities:
-            return False
-
-        new = {}
-        for e in entities:
-            try:
-                p = utils.get_input_peer(e)
-                new[utils.get_peer_id(p, add_mark=True)] = \
-                    getattr(p, 'access_hash', 0)  # chats won't have hash
-            except ValueError:
-                pass
-
-        with self._entities_lock:
-            before = len(self._input_entities)
-            self._input_entities.update(new)
-            return len(self._input_entities) != before
-
-    def get_input_entity(self, peer):
-        """Gets an input entity known its Peer or a marked ID,
-           or raises KeyError if not found/invalid.
-        """
-        if not isinstance(peer, int):
-            peer = utils.get_peer_id(peer, add_mark=True)
-
-        entity_hash = self._input_entities[peer]
-        entity_id, peer_class = utils.resolve_id(peer)
-
-        if peer_class == PeerUser:
-            return InputPeerUser(entity_id, entity_hash)
-        if peer_class == PeerChat:
-            return InputPeerChat(entity_id)
-        if peer_class == PeerChannel:
-            return InputPeerChannel(entity_id, entity_hash)
-
-        raise KeyError()

From e5c4df98df4d7f003359bf0d3419788c8bd49fa8 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 5 Oct 2017 12:27:05 +0200
Subject: [PATCH 03/11] Use EntityDatabase on TelegramClient.get_entity instead
 lru_cache

---
 telethon/telegram_client.py    | 111 +++++++++++++++++----------------
 telethon/tl/entity_database.py |   6 ++
 2 files changed, 62 insertions(+), 55 deletions(-)

diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py
index 5fb81cc6..138302aa 100644
--- a/telethon/telegram_client.py
+++ b/telethon/telegram_client.py
@@ -849,7 +849,6 @@ class TelegramClient(TelegramBareClient):
 
     # region Small utilities to make users' life easier
 
-    @lru_cache()
     def get_entity(self, entity):
         """Turns an entity into a valid Telegram user or chat.
            If "entity" is a string which can be converted to an integer,
@@ -866,61 +865,62 @@ class TelegramClient(TelegramBareClient):
            If the entity is neither, and it's not a TLObject, an
            error will be raised.
         """
-        # TODO Maybe cache both the contacts and the entities.
-        # If an user cannot be found, force a cache update through
-        # a public method (since users may change their username)
-        input_entity = None
-        if isinstance(entity, TLObject):
-            # crc32(b'InputPeer') and crc32(b'Peer')
-            if type(entity).SUBCLASS_OF_ID in (0xc91c90b6, 0x2d45687):
-                input_entity = self.get_input_entity(entity)
-            else:
-                # TODO Don't assume it's a valid entity
-                return entity
+        try:
+            return self.session.entities[entity]
+        except KeyError:
+            pass
 
-        elif isinstance(entity, int):
-            input_entity = self.get_input_entity(entity)
+        if isinstance(entity, int) or (
+                isinstance(entity, TLObject) and
+                # crc32(b'InputPeer') and crc32(b'Peer')
+                type(entity).SUBCLASS_OF_ID in (0xc91c90b6, 0x2d45687)):
+            ie = self.get_input_entity(entity)
+            if isinstance(ie, InputPeerUser):
+                self.session.process_entities(GetUsersRequest([ie]))
+            elif isinstance(ie, InputPeerChat):
+                self.session.process_entities(GetChatsRequest([ie.chat_id]))
+            elif isinstance(ie, InputPeerChannel):
+                self.session.process_entities(GetChannelsRequest([ie]))
 
-        if input_entity:
-            if isinstance(input_entity, InputPeerUser):
-                return self(GetUsersRequest([input_entity]))[0]
-            elif isinstance(input_entity, InputPeerChat):
-                return self(GetChatsRequest([input_entity.chat_id])).chats[0]
-            elif isinstance(input_entity, InputPeerChannel):
-                return self(GetChannelsRequest([input_entity])).chats[0]
+            # The result of Get*Request has been processed and the entity
+            # cached if it was found.
+            return self.session.entities[ie]
 
         if isinstance(entity, str):
-            stripped_phone = self._parse_phone(entity, ignore_saved=True)
-            if stripped_phone.isdigit():
-                contacts = self(GetContactsRequest(0))
-                try:
-                    return next(
-                        u for u in contacts.users
-                        if u.phone and u.phone.endswith(stripped_phone)
-                    )
-                except StopIteration:
-                    raise ValueError(
-                        'Could not find user with phone {}, '
-                        'add them to your contacts first'.format(entity)
-                    )
-            else:
-                username = entity.strip('@').lower()
-                resolved = self(ResolveUsernameRequest(username))
-                for c in resolved.chats:
-                    if getattr(c, 'username', '').lower() == username:
-                        return c
-                for u in resolved.users:
-                    if getattr(u, 'username', '').lower() == username:
-                        return u
-
-                raise ValueError(
-                    'Could not find user with username {}'.format(entity)
-                )
+            return self._get_entity_from_string(entity)
 
         raise ValueError(
             'Cannot turn "{}" into any entity (user or chat)'.format(entity)
         )
 
+    def _get_entity_from_string(self, string):
+        """Gets an entity from the given string, which may be a phone or
+           an username, and processes all the found entities on the session.
+        """
+        stripped_phone = self._parse_phone(string, ignore_saved=True)
+        if stripped_phone.isdigit():
+            contacts = self(GetContactsRequest(0))
+            self.session.process_entities(contacts)
+            try:
+                return next(
+                    u for u in contacts.users
+                    if u.phone and u.phone.endswith(stripped_phone)
+                )
+            except StopIteration:
+                raise ValueError(
+                    'Could not find user with phone {}, '
+                    'add them to your contacts first'.format(string)
+                )
+        else:
+            entity = string.strip('@').lower()
+            self.session.process_entities(self(ResolveUsernameRequest(entity)))
+            try:
+                return self.session.entities[entity]
+            except KeyError:
+                raise ValueError(
+                    'Could not find user with username {}'.format(entity)
+                )
+
     def _parse_phone(self, phone, ignore_saved=False):
         if isinstance(phone, int):
             phone = str(phone)
@@ -942,9 +942,14 @@ class TelegramClient(TelegramBareClient):
 
            If even after
         """
+        try:
+            # First try to get the entity from cache, otherwise figure it out
+            self.session.entities.get_input_entity(peer)
+        except KeyError:
+            pass
+
         if isinstance(peer, str):
-            # Let .get_entity resolve the username or phone (full entity)
-            peer = self.get_entity(peer)
+            return utils.get_input_peer(self._get_entity_from_string(peer))
 
         is_peer = False
         if isinstance(peer, int):
@@ -964,16 +969,12 @@ class TelegramClient(TelegramBareClient):
                 'Cannot turn "{}" into an input entity.'.format(peer)
             )
 
-        try:
-            return self.session.get_input_entity(peer)
-        except KeyError:
-            pass
-
         if self.session.save_entities:
             # Not found, look in the dialogs (this will save the users)
             self.get_dialogs(limit=None)
+
             try:
-                return self.session.get_input_entity(peer)
+                self.session.entities.get_input_entity(peer)
             except KeyError:
                 pass
 
diff --git a/telethon/tl/entity_database.py b/telethon/tl/entity_database.py
index 6a1d7dbb..50a4207e 100644
--- a/telethon/tl/entity_database.py
+++ b/telethon/tl/entity_database.py
@@ -130,6 +130,12 @@ class EntityDatabase:
 
     # TODO Allow search by name by tokenizing the input and return a list
 
+    def get_input_entity(self, peer):
+        try:
+            return self._input_entities[utils.get_peer_id(peer, add_mark=True)]
+        except ValueError as e:
+            raise KeyError(peer) from e
+
     def get_input_list(self):
         return list(self._input_entities.items())
 

From d2e244817aa4d5ac21469038cd529e3f3163365d Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 5 Oct 2017 12:28:04 +0200
Subject: [PATCH 04/11] Make EntityDatabase.get a no-op if key is an entity
 already

---
 telethon/tl/entity_database.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/telethon/tl/entity_database.py b/telethon/tl/entity_database.py
index 50a4207e..94528cfd 100644
--- a/telethon/tl/entity_database.py
+++ b/telethon/tl/entity_database.py
@@ -117,8 +117,14 @@ class EntityDatabase:
         if isinstance(key, int):
             return self._entities[key]  # normal IDs are assumed users
 
-        if isinstance(key, TLObject) and type(key).SUBCLASS_OF_ID == 0x2d45687:
-            return self._entities[utils.get_peer_id(key, add_mark=True)]
+        if isinstance(key, TLObject):
+            sc = type(key).SUBCLASS_OF_ID
+            if sc == 0x2d45687:
+                # Subclass of "Peer"
+                return self._entities[utils.get_peer_id(key, add_mark=True)]
+            elif sc in {0x2da17977, 0xc5af5d94, 0x6d44b7db}:
+                # Subclass of "User", "Chat" or "Channel"
+                return key
 
         raise KeyError(key)
 

From 10eca821439cf4f2c2858bcb6585bcd75c404c2b Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 5 Oct 2017 12:29:52 +0200
Subject: [PATCH 05/11] Fix process_entities not working on list of entities

---
 telethon/tl/entity_database.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/telethon/tl/entity_database.py b/telethon/tl/entity_database.py
index 94528cfd..93a9bb71 100644
--- a/telethon/tl/entity_database.py
+++ b/telethon/tl/entity_database.py
@@ -30,6 +30,10 @@ class EntityDatabase:
             return False
 
         # Save all input entities we know of
+        if not isinstance(tlobject, TLObject) and hasattr(tlobject, '__iter__'):
+            # This may be a list of users already for instance
+            return self.expand(tlobject)
+
         entities = []
         if hasattr(tlobject, 'chats') and hasattr(tlobject.chats, '__iter__'):
             entities.extend(tlobject.chats)

From 16f929b8b65ff38f53183847c8e5a10441928b3b Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 5 Oct 2017 12:33:24 +0200
Subject: [PATCH 06/11] Fix .get_peer_id not working with full entities

---
 telethon/utils.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/telethon/utils.py b/telethon/utils.py
index f6e6373a..8f213054 100644
--- a/telethon/utils.py
+++ b/telethon/utils.py
@@ -321,7 +321,11 @@ def get_peer_id(peer, add_mark=False):
             i = peer.channel_id  # IDs will be strictly positive -> log works
             return -(i + pow(10, math.floor(math.log10(i) + 3)))
 
-    _raise_cast_fail(peer, 'int')
+    # Maybe a full entity was given and we just need its ID
+    try:
+        return get_peer_id(get_input_peer(peer), add_mark=add_mark)
+    except ValueError:
+        _raise_cast_fail(peer, 'int')
 
 
 def resolve_id(marked_id):

From a8edacd34a7f5e7842fb2187b9de4beb8cf072ea Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 5 Oct 2017 12:59:44 +0200
Subject: [PATCH 07/11] Fix get_peer_id going into infinite recursion for
 InputPeerSelf

---
 telethon/utils.py | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/telethon/utils.py b/telethon/utils.py
index 8f213054..e021af03 100644
--- a/telethon/utils.py
+++ b/telethon/utils.py
@@ -304,6 +304,16 @@ def get_peer_id(peer, add_mark=False):
     """Finds the ID of the given peer, and optionally converts it to
        the "bot api" format if 'add_mark' is set to True.
     """
+    if not isinstance(peer, TLObject):
+        if isinstance(peer, int):
+            return peer
+        else:
+            _raise_cast_fail(peer, 'int')
+
+    if type(peer).SUBCLASS_OF_ID not in {0x2d45687, 0xc91c90b6}:
+        # Not a Peer or an InputPeer, so first get its Input version
+        peer = get_input_peer(peer)
+
     if isinstance(peer, PeerUser) or isinstance(peer, InputPeerUser):
         return peer.user_id
     else:
@@ -321,11 +331,7 @@ def get_peer_id(peer, add_mark=False):
             i = peer.channel_id  # IDs will be strictly positive -> log works
             return -(i + pow(10, math.floor(math.log10(i) + 3)))
 
-    # Maybe a full entity was given and we just need its ID
-    try:
-        return get_peer_id(get_input_peer(peer), add_mark=add_mark)
-    except ValueError:
-        _raise_cast_fail(peer, 'int')
+    _raise_cast_fail(peer, 'int')
 
 
 def resolve_id(marked_id):

From 99cc0778bb20241df353df7259ce44c06f64b528 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 5 Oct 2017 13:01:00 +0200
Subject: [PATCH 08/11] Allow EntityDatabase to be accessed by phone

---
 telethon/telegram_client.py    | 49 ++++++++++------------------------
 telethon/tl/entity_database.py | 35 +++++++++++++++++++-----
 2 files changed, 43 insertions(+), 41 deletions(-)

diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py
index 138302aa..5d8a2dd6 100644
--- a/telethon/telegram_client.py
+++ b/telethon/telegram_client.py
@@ -1,7 +1,5 @@
 import os
-import re
 from datetime import datetime, timedelta
-from functools import lru_cache
 from mimetypes import guess_type
 
 try:
@@ -17,6 +15,7 @@ from .errors import (
 )
 from .network import ConnectionMode
 from .tl import TLObject
+from .tl.entity_database import EntityDatabase
 from .tl.functions.account import (
     GetPasswordRequest
 )
@@ -127,7 +126,7 @@ class TelegramClient(TelegramBareClient):
 
     def send_code_request(self, phone):
         """Sends a code request to the specified phone number"""
-        phone = self._parse_phone(phone)
+        phone = EntityDatabase.parse_phone(phone) or self._phone
         result = self(SendCodeRequest(phone, self.api_id, self.api_hash))
         self._phone = phone
         self._phone_code_hash = result.phone_code_hash
@@ -160,7 +159,7 @@ class TelegramClient(TelegramBareClient):
         if phone and not code:
             return self.send_code_request(phone)
         elif code:
-            phone = self._parse_phone(phone)
+            phone = EntityDatabase.parse_phone(phone) or self._phone
             phone_code_hash = phone_code_hash or self._phone_code_hash
             if not phone:
                 raise ValueError(
@@ -897,40 +896,20 @@ class TelegramClient(TelegramBareClient):
         """Gets an entity from the given string, which may be a phone or
            an username, and processes all the found entities on the session.
         """
-        stripped_phone = self._parse_phone(string, ignore_saved=True)
-        if stripped_phone.isdigit():
-            contacts = self(GetContactsRequest(0))
-            self.session.process_entities(contacts)
-            try:
-                return next(
-                    u for u in contacts.users
-                    if u.phone and u.phone.endswith(stripped_phone)
-                )
-            except StopIteration:
-                raise ValueError(
-                    'Could not find user with phone {}, '
-                    'add them to your contacts first'.format(string)
-                )
+        phone = EntityDatabase.parse_phone(string)
+        if phone:
+            entity = phone
+            self.session.process_entities(self(GetContactsRequest(0)))
         else:
             entity = string.strip('@').lower()
             self.session.process_entities(self(ResolveUsernameRequest(entity)))
-            try:
-                return self.session.entities[entity]
-            except KeyError:
-                raise ValueError(
-                    'Could not find user with username {}'.format(entity)
-                )
 
-    def _parse_phone(self, phone, ignore_saved=False):
-        if isinstance(phone, int):
-            phone = str(phone)
-        elif phone:
-            phone = re.sub(r'[+()\s-]', '', phone)
-
-        if ignore_saved:
-            return phone
-        else:
-            return phone or self._phone
+        try:
+            return self.session.entities[entity]
+        except KeyError:
+            raise ValueError(
+                'Could not find user with username {}'.format(entity)
+            )
 
     def get_input_entity(self, peer):
         """Gets the input entity given its PeerUser, PeerChat, PeerChannel.
@@ -940,7 +919,7 @@ class TelegramClient(TelegramBareClient):
            If this Peer hasn't been seen before by the library, all dialogs
            will loaded, and their entities saved to the session file.
 
-           If even after
+           If even after it's not found, a ValueError is raised.
         """
         try:
             # First try to get the entity from cache, otherwise figure it out
diff --git a/telethon/tl/entity_database.py b/telethon/tl/entity_database.py
index 93a9bb71..b42aec3c 100644
--- a/telethon/tl/entity_database.py
+++ b/telethon/tl/entity_database.py
@@ -1,5 +1,7 @@
 from threading import Lock
 
+import re
+
 from .. import utils
 from ..tl import TLObject
 from ..tl.types import User, Chat, Channel
@@ -19,6 +21,7 @@ class EntityDatabase:
 
         # TODO Allow disabling some extra mappings
         self._username_id = {}  # username: marked_id
+        self._phone_id = {}  # phone: marked_id
 
     def process(self, tlobject):
         """Processes all the found entities on the given TLObject,
@@ -87,19 +90,27 @@ class EntityDatabase:
             old_entity = self._entities[marked_id]
             old_entity.__dict__.update(entity.__dict__)  # Keep old references
 
-            # Update must delete old username
+            # Update must delete old username and phone
             username = getattr(old_entity, 'username', None)
             if username:
                 del self._username_id[username.lower()]
+
+            phone = getattr(old_entity, 'phone', None)
+            if phone:
+                del self._phone_id[phone]
         except KeyError:
             # Add new entity
             self._entities[marked_id] = entity
 
-        # Always update username if any
+        # Always update username or phone if any
         username = getattr(entity, 'username', None)
         if username:
             self._username_id[username.lower()] = marked_id
 
+        phone = getattr(entity, 'phone', None)
+        if phone:
+            self._username_id[phone] = marked_id
+
     def __getitem__(self, key):
         """Accepts a digit only string as phone number,
            otherwise it's treated as an username.
@@ -113,10 +124,12 @@ class EntityDatabase:
            Note that megagroups are channels with .megagroup = True.
         """
         if isinstance(key, str):
-            # TODO Parse phone properly, currently only usernames
-            key = key.lstrip('@').lower()
-            # TODO Use the client to return from username if not found
-            return self._entities[self._username_id[key]]
+            phone = EntityDatabase.parse_phone(key)
+            if phone:
+                return self._phone_id[phone]
+            else:
+                key = key.lstrip('@').lower()
+                return self._entities[self._username_id[key]]
 
         if isinstance(key, int):
             return self._entities[key]  # normal IDs are assumed users
@@ -140,6 +153,16 @@ class EntityDatabase:
 
     # TODO Allow search by name by tokenizing the input and return a list
 
+    @staticmethod
+    def parse_phone(phone):
+        """Parses the given phone, or returns None if it's invalid"""
+        if isinstance(phone, int):
+            return str(phone)
+        else:
+            phone = re.sub(r'[+()\s-]', '', str(phone))
+            if phone.isdigit():
+                return phone
+
     def get_input_entity(self, peer):
         try:
             return self._input_entities[utils.get_peer_id(peer, add_mark=True)]

From 4f2a44231a32e96f25f41f3f01af4f14dc8319c3 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 5 Oct 2017 13:06:51 +0200
Subject: [PATCH 09/11] Allow disabling the EntityDatabase fully or partially

---
 telethon/tl/entity_database.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/telethon/tl/entity_database.py b/telethon/tl/entity_database.py
index b42aec3c..ee2c5b67 100644
--- a/telethon/tl/entity_database.py
+++ b/telethon/tl/entity_database.py
@@ -8,8 +8,15 @@ from ..tl.types import User, Chat, Channel
 
 
 class EntityDatabase:
-    def __init__(self, input_list=None, enabled=True):
+    def __init__(self, input_list=None, enabled=True, enabled_full=True):
+        """Creates a new entity database with an initial load of "Input"
+           entities, if any.
+
+           If 'enabled', input entities will be saved. The whole entity
+           will be saved if both 'enabled' and 'enabled_full' are True.
+        """
         self.enabled = enabled
+        self.enabled_full = enabled_full
 
         self._lock = Lock()
         self._entities = {}  # marked_id: user|chat|channel
@@ -63,10 +70,11 @@ class EntityDatabase:
                 new_input[utils.get_peer_id(p, add_mark=True)] = \
                     getattr(p, 'access_hash', 0)  # chats won't have hash
 
-                if isinstance(e, User) \
-                        or isinstance(e, Chat) \
-                        or isinstance(e, Channel):
-                    new.append(e)
+                if self.enabled_full:
+                    if isinstance(e, User) \
+                            or isinstance(e, Chat) \
+                            or isinstance(e, Channel):
+                        new.append(e)
             except ValueError:
                 pass
 
@@ -78,7 +86,8 @@ class EntityDatabase:
             return len(self._input_entities) != before
 
     def _add_full_entity(self, entity):
-        """Adds a "full" entity (User, Chat or Channel, not "Input*").
+        """Adds a "full" entity (User, Chat or Channel, not "Input*"),
+           despite the value of self.enabled and self.enabled_full.
 
            Not to be confused with UserFull, ChatFull, or ChannelFull,
            "full" means simply not "Input*".

From 1fb3d0d00cf8c778b7b483ba54614ee360f9946a Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 5 Oct 2017 13:14:54 +0200
Subject: [PATCH 10/11] Fix EntityDatabase failing to cache self user

---
 telethon/tl/entity_database.py | 4 ++--
 telethon/utils.py              | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/telethon/tl/entity_database.py b/telethon/tl/entity_database.py
index ee2c5b67..ce25a6e0 100644
--- a/telethon/tl/entity_database.py
+++ b/telethon/tl/entity_database.py
@@ -66,7 +66,7 @@ class EntityDatabase:
                 continue
 
             try:
-                p = utils.get_input_peer(e)
+                p = utils.get_input_peer(e, allow_self=False)
                 new_input[utils.get_peer_id(p, add_mark=True)] = \
                     getattr(p, 'access_hash', 0)  # chats won't have hash
 
@@ -93,7 +93,7 @@ class EntityDatabase:
            "full" means simply not "Input*".
         """
         marked_id = utils.get_peer_id(
-            utils.get_input_peer(entity), add_mark=True
+            utils.get_input_peer(entity, allow_self=False), add_mark=True
         )
         try:
             old_entity = self._entities[marked_id]
diff --git a/telethon/utils.py b/telethon/utils.py
index e021af03..74fbb90d 100644
--- a/telethon/utils.py
+++ b/telethon/utils.py
@@ -72,7 +72,7 @@ def _raise_cast_fail(entity, target):
                      .format(type(entity).__name__, target))
 
 
-def get_input_peer(entity):
+def get_input_peer(entity, allow_self=True):
     """Gets the input peer for the given "entity" (user, chat or channel).
        A ValueError is raised if the given entity isn't a supported type."""
     if not isinstance(entity, TLObject):
@@ -82,7 +82,7 @@ def get_input_peer(entity):
         return entity
 
     if isinstance(entity, User):
-        if entity.is_self:
+        if entity.is_self and allow_self:
             return InputPeerSelf()
         else:
             return InputPeerUser(entity.id, entity.access_hash)
@@ -312,7 +312,7 @@ def get_peer_id(peer, add_mark=False):
 
     if type(peer).SUBCLASS_OF_ID not in {0x2d45687, 0xc91c90b6}:
         # Not a Peer or an InputPeer, so first get its Input version
-        peer = get_input_peer(peer)
+        peer = get_input_peer(peer, allow_self=False)
 
     if isinstance(peer, PeerUser) or isinstance(peer, InputPeerUser):
         return peer.user_id

From dde196d8e4ac46331a3f646b5eaf84b38770175b Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 5 Oct 2017 13:34:44 +0200
Subject: [PATCH 11/11] Swallow errors when processing entities

Some bug was causing an infinite lock due to the entity
database failing with InputPeerSelf, since adding a full
entity to the database wasn't disallowing self, and so
wasn't utils get_peer_id. Although last commit fixed that,
just in case, swallow errors there.
---
 telethon/tl/session.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/telethon/tl/session.py b/telethon/tl/session.py
index 2b691ad7..98ffda16 100644
--- a/telethon/tl/session.py
+++ b/telethon/tl/session.py
@@ -181,5 +181,8 @@ class Session:
         self.time_offset = correct - now
 
     def process_entities(self, tlobject):
-        if self.entities.process(tlobject):
-            self.save()  # Save if any new entities got added
+        try:
+            if self.entities.process(tlobject):
+                self.save()  # Save if any new entities got added
+        except:
+            pass