From 8abc7ade224620af76ebad1896424eeb5ae65fc7 Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Tue, 26 Mar 2019 11:39:25 +0100 Subject: [PATCH] Use the new in-memory entity cache This should avoid a disk access every time an input entity is needed, which is very often. Another step for #1141. --- telethon/client/telegrambaseclient.py | 1 - telethon/client/users.py | 4 ++-- telethon/entitycache.py | 10 ++++++---- telethon/events/callbackquery.py | 6 ++---- telethon/events/chataction.py | 6 ++---- telethon/events/common.py | 6 ++---- telethon/events/userupdate.py | 4 ++-- telethon/tl/custom/chatgetter.py | 5 ++--- telethon/tl/custom/draft.py | 5 ++--- telethon/tl/custom/message.py | 6 ++++-- 10 files changed, 24 insertions(+), 29 deletions(-) diff --git a/telethon/client/telegrambaseclient.py b/telethon/client/telegrambaseclient.py index 56325312..ab21cae0 100644 --- a/telethon/client/telegrambaseclient.py +++ b/telethon/client/telegrambaseclient.py @@ -235,7 +235,6 @@ class TelegramBaseClient(abc.ABC): # # The session files only wants the entities to persist # them to disk, and to save additional useful information. - # TODO Make use of _entity_cache # TODO Session should probably return all cached # info of entities, not just the input versions self.session = session diff --git a/telethon/client/users.py b/telethon/client/users.py index 6bbf86ab..28b1df55 100644 --- a/telethon/client/users.py +++ b/telethon/client/users.py @@ -308,8 +308,8 @@ class UserMethods(TelegramBaseClient): try: # First try to get the entity from cache, otherwise figure it out - return self.session.get_input_entity(peer) - except ValueError: + return self._entity_cache[peer] + except KeyError: pass if isinstance(peer, str): diff --git a/telethon/entitycache.py b/telethon/entitycache.py index e74af793..0f94ae32 100644 --- a/telethon/entitycache.py +++ b/telethon/entitycache.py @@ -31,15 +31,17 @@ class EntityCache: def __getitem__(self, item): """ Gets the corresponding :tl:`InputPeer` for the given ID or peer, - or returns `None` on error/not found. + or raises ``KeyError`` on any error (i.e. cannot be found). """ if not isinstance(item, int) or item < 0: try: - return self.__dict__.get(utils.get_peer_id(item)) + return self.__dict__[utils.get_peer_id(item)] except TypeError: - return None + raise KeyError('Invalid key will not have entity') from None for cls in (types.PeerUser, types.PeerChat, types.PeerChannel): - result = self.__dict__.get(cls(item)) + result = self.__dict__.get(utils.get_peer_id(cls(item))) if result: return result + + raise KeyError('No cached entity for the given key') diff --git a/telethon/events/callbackquery.py b/telethon/events/callbackquery.py index 527a4e84..975800e3 100644 --- a/telethon/events/callbackquery.py +++ b/telethon/events/callbackquery.py @@ -156,10 +156,8 @@ class CallbackQuery(EventBuilder): if not getattr(self._input_sender, 'access_hash', True): # getattr with True to handle the InputPeerSelf() case try: - self._input_sender = self._client.session.get_input_entity( - self._sender_id - ) - except ValueError: + self._input_sender = self._client._entity_cache[self._sender_id] + except KeyError: m = await self.get_message() if m: self._sender = m._sender diff --git a/telethon/events/chataction.py b/telethon/events/chataction.py index 52d6b3ab..4ed54519 100644 --- a/telethon/events/chataction.py +++ b/telethon/events/chataction.py @@ -362,10 +362,8 @@ class ChatAction(EventBuilder): self._input_users = [] for peer in self._user_peers: try: - self._input_users.append( - self._client.session.get_input_entity(peer) - ) - except ValueError: + self._input_users.append(self._client._entity_cache[peer]) + except KeyError: pass return self._input_users or [] diff --git a/telethon/events/common.py b/telethon/events/common.py index 8b9bdf92..1470585a 100644 --- a/telethon/events/common.py +++ b/telethon/events/common.py @@ -154,10 +154,8 @@ class EventCommon(ChatGetter, abc.ABC): self._input_chat = utils.get_input_peer(self._chat) except TypeError: try: - self._input_chat = self._client.session.get_input_entity( - self._chat_peer - ) - except ValueError: + self._input_chat = self._client._entity_cache[self._chat_peer] + except KeyError: self._input_chat = None @property diff --git a/telethon/events/userupdate.py b/telethon/events/userupdate.py index a4fdf533..9d4dab44 100644 --- a/telethon/events/userupdate.py +++ b/telethon/events/userupdate.py @@ -170,7 +170,7 @@ class UserUpdate(EventBuilder): def _set_client(self, client): if isinstance(self._chat_peer, int): try: - chat = client.session.get_input_entity(self._chat_peer) + chat = client._entity_cache[self._chat_peer] if isinstance(chat, types.InputPeerChat): self._chat_peer = types.PeerChat(self._chat_peer) elif isinstance(chat, types.InputPeerChannel): @@ -178,7 +178,7 @@ class UserUpdate(EventBuilder): else: # Should not happen self._chat_peer = types.PeerUser(self._chat_peer) - except ValueError: + except KeyError: # Hope for the best. We don't know where this event # occurred but it was most likely in a channel. self._chat_peer = types.PeerChannel(self._chat_peer) diff --git a/telethon/tl/custom/chatgetter.py b/telethon/tl/custom/chatgetter.py index fb733d9a..8eb904d4 100644 --- a/telethon/tl/custom/chatgetter.py +++ b/telethon/tl/custom/chatgetter.py @@ -49,9 +49,8 @@ class ChatGetter(abc.ABC): """ if self._input_chat is None and self._chat_peer: try: - self._input_chat =\ - self._client.session.get_input_entity(self._chat_peer) - except ValueError: + self._input_chat = self._client._entity_cache(self._chat_peer) + except KeyError: pass return self._input_chat diff --git a/telethon/tl/custom/draft.py b/telethon/tl/custom/draft.py index de5e5991..eb12ad3d 100644 --- a/telethon/tl/custom/draft.py +++ b/telethon/tl/custom/draft.py @@ -64,9 +64,8 @@ class Draft: """ if not self._input_entity: try: - self._input_entity =\ - self._client.session.get_input_entity(self._peer) - except ValueError: + self._input_entity = self._client._entity_cache[self._peer] + except KeyError: pass return self._input_entity diff --git a/telethon/tl/custom/message.py b/telethon/tl/custom/message.py index d8c2b4ca..9fcdda70 100644 --- a/telethon/tl/custom/message.py +++ b/telethon/tl/custom/message.py @@ -834,8 +834,10 @@ class Message(ChatGetter, SenderGetter, TLObject, abc.ABC): if not bot: raise ValueError('No input sender') else: - return self._client.session.get_input_entity( - self.via_bot_id) + try: + return self._client._entity_cache[self.via_bot_id] + except KeyError: + raise ValueError('No input sender') from None def _document_by_attribute(self, kind, condition=None): """