Use EntityDatabase on TelegramClient.get_entity instead lru_cache

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.
+        try:
-        # If an user cannot be found, force a cache update through
+            return self.session.entities[entity]
-        # a public method (since users may change their username)
+        except KeyError:
-        input_entity = None
+            pass
-        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
 
-        elif isinstance(entity, int):
+        if isinstance(entity, int) or (
-            input_entity = self.get_input_entity(entity)
+                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:
+            # The result of Get*Request has been processed and the entity
-            if isinstance(input_entity, InputPeerUser):
+            # cached if it was found.
-                return self(GetUsersRequest([input_entity]))[0]
+            return self.session.entities[ie]
-            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]
 
         if isinstance(entity, str):
-            stripped_phone = self._parse_phone(entity, ignore_saved=True)
+            return self._get_entity_from_string(entity)
-            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)
-                )
 
         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)
+            return utils.get_input_peer(self._get_entity_from_string(peer))
-            peer = self.get_entity(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
 

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())