From e5c4df98df4d7f003359bf0d3419788c8bd49fa8 Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Thu, 5 Oct 2017 12:27:05 +0200 Subject: [PATCH] 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())