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,33 +865,42 @@ 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):
+        try:
+            return self.session.entities[entity]
+        except KeyError:
+            pass
+
+        if isinstance(entity, int) or (
+                isinstance(entity, TLObject) and
                 # 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
+                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]))
 
-        elif isinstance(entity, int):
-            input_entity = self.get_input_entity(entity)
-
-        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)
+            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
@@ -901,26 +909,18 @@ class TelegramClient(TelegramBareClient):
             except StopIteration:
                 raise ValueError(
                     'Could not find user with phone {}, '
-                        'add them to your contacts first'.format(entity)
+                    'add them to your contacts first'.format(string)
                 )
         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
-
+            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)
                 )
 
-        raise ValueError(
-            'Cannot turn "{}" into any entity (user or chat)'.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
 

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