diff --git a/telethon/client/telegrambaseclient.py b/telethon/client/telegrambaseclient.py index 54c02a16..d764c99c 100644 --- a/telethon/client/telegrambaseclient.py +++ b/telethon/client/telegrambaseclient.py @@ -300,7 +300,7 @@ class TelegramBaseClient(abc.ABC): self.flood_sleep_threshold = flood_sleep_threshold # TODO Use AsyncClassWrapper(session) - # ChatGetter and SenderGetter can use the in-memory _entity_cache + # ChatGetter and SenderGetter can use the in-memory _mb_entity_cache # to avoid network access and the need for await in session files. # # The session files only wants the entities to persist @@ -308,7 +308,6 @@ class TelegramBaseClient(abc.ABC): # TODO Session should probably return all cached # info of entities, not just the input versions self.session = session - self._entity_cache = EntityCache() self.api_id = int(api_id) self.api_hash = api_hash @@ -433,7 +432,6 @@ class TelegramBaseClient(abc.ABC): self._catch_up = catch_up self._updates_queue = asyncio.Queue() self._message_box = MessageBox(self._log['messagebox']) - # This entity cache is tailored for the messagebox and is not used for absolutely everything like _entity_cache self._mb_entity_cache = MbEntityCache() # required for proper update handling (to know when to getDifference) self._sender = MTProtoSender( diff --git a/telethon/client/users.py b/telethon/client/users.py index 85f8ff03..50f9ff81 100644 --- a/telethon/client/users.py +++ b/telethon/client/users.py @@ -72,7 +72,6 @@ class UserMethods: results.append(None) continue self.session.process_entities(result) - self._entity_cache.add(result) exceptions.append(None) results.append(result) request_index += 1 @@ -83,7 +82,6 @@ class UserMethods: else: result = await future self.session.process_entities(result) - self._entity_cache.add(result) return result except (errors.ServerError, errors.RpcCallFailError, errors.RpcMcgetFailError, errors.InterdcCallErrorError, @@ -417,8 +415,8 @@ class UserMethods: try: # 0x2d45687 == crc32(b'Peer') if isinstance(peer, int) or peer.SUBCLASS_OF_ID == 0x2d45687: - return self._entity_cache[peer] - except (AttributeError, KeyError): + return self._mb_entity_cache.get(utils.get_peer_id(peer, add_mark=False))._as_input_peer() + except AttributeError: pass # Then come known strings that take precedence diff --git a/telethon/events/album.py b/telethon/events/album.py index 827f1b11..481fedf4 100644 --- a/telethon/events/album.py +++ b/telethon/events/album.py @@ -160,7 +160,7 @@ class Album(EventBuilder): def _set_client(self, client): super()._set_client(client) self._sender, self._input_sender = utils._get_entity_pair( - self.sender_id, self._entities, client._entity_cache) + self.sender_id, self._entities, client._mb_entity_cache) for msg in self.messages: msg._finish_init(client, self._entities, None) diff --git a/telethon/events/callbackquery.py b/telethon/events/callbackquery.py index 94e03b7b..ce8f457c 100644 --- a/telethon/events/callbackquery.py +++ b/telethon/events/callbackquery.py @@ -151,7 +151,7 @@ class CallbackQuery(EventBuilder): def _set_client(self, client): super()._set_client(client) self._sender, self._input_sender = utils._get_entity_pair( - self.sender_id, self._entities, client._entity_cache) + self.sender_id, self._entities, client._mb_entity_cache) @property def id(self): @@ -208,8 +208,9 @@ 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._entity_cache[self._sender_id] - except KeyError: + self._input_sender = self._client._mb_entity_cache.get( + utils.resolve_id(self._sender_id)[0])._as_input_peer() + except AttributeError: m = await self.get_message() if m: self._sender = m._sender diff --git a/telethon/events/chataction.py b/telethon/events/chataction.py index 46563059..b8fc3dc2 100644 --- a/telethon/events/chataction.py +++ b/telethon/events/chataction.py @@ -425,9 +425,10 @@ class ChatAction(EventBuilder): # If missing, try from the entity cache try: - self._input_users.append(self._client._entity_cache[user_id]) + self._input_users.append(self._client._mb_entity_cache.get( + utils.resolve_id(user_id)[0])._as_input_peer()) continue - except KeyError: + except AttributeError: pass return self._input_users or [] diff --git a/telethon/events/common.py b/telethon/events/common.py index e8978f31..e295adb2 100644 --- a/telethon/events/common.py +++ b/telethon/events/common.py @@ -154,7 +154,7 @@ class EventCommon(ChatGetter, abc.ABC): self._client = client if self._chat_peer: self._chat, self._input_chat = utils._get_entity_pair( - self.chat_id, self._entities, client._entity_cache) + self.chat_id, self._entities, client._mb_entity_cache) else: self._chat = self._input_chat = None diff --git a/telethon/events/inlinequery.py b/telethon/events/inlinequery.py index 98d8e124..7a065e55 100644 --- a/telethon/events/inlinequery.py +++ b/telethon/events/inlinequery.py @@ -99,7 +99,7 @@ class InlineQuery(EventBuilder): def _set_client(self, client): super()._set_client(client) self._sender, self._input_sender = utils._get_entity_pair( - self.sender_id, self._entities, client._entity_cache) + self.sender_id, self._entities, client._mb_entity_cache) @property def id(self): diff --git a/telethon/events/userupdate.py b/telethon/events/userupdate.py index 1516f8ee..6417f2b4 100644 --- a/telethon/events/userupdate.py +++ b/telethon/events/userupdate.py @@ -95,7 +95,7 @@ class UserUpdate(EventBuilder): def _set_client(self, client): super()._set_client(client) self._sender, self._input_sender = utils._get_entity_pair( - self.sender_id, self._entities, client._entity_cache) + self.sender_id, self._entities, client._mb_entity_cache) @property def user(self): diff --git a/telethon/tl/custom/chatgetter.py b/telethon/tl/custom/chatgetter.py index 60f5c79e..240d9b83 100644 --- a/telethon/tl/custom/chatgetter.py +++ b/telethon/tl/custom/chatgetter.py @@ -66,8 +66,9 @@ class ChatGetter(abc.ABC): """ if self._input_chat is None and self._chat_peer and self._client: try: - self._input_chat = self._client._entity_cache[self._chat_peer] - except KeyError: + self._input_chat = self._client._mb_entity_cache.get( + utils.get_peer_id(self._chat_peer, add_mark=False))._as_input_peer() + except AttributeError: pass return self._input_chat diff --git a/telethon/tl/custom/draft.py b/telethon/tl/custom/draft.py index a44986b7..c482b64f 100644 --- a/telethon/tl/custom/draft.py +++ b/telethon/tl/custom/draft.py @@ -5,7 +5,7 @@ from ..functions.messages import SaveDraftRequest from ..types import DraftMessage from ...errors import RPCError from ...extensions import markdown -from ...utils import get_input_peer, get_peer +from ...utils import get_input_peer, get_peer, get_peer_id class Draft: @@ -53,8 +53,9 @@ class Draft: """ if not self._input_entity: try: - self._input_entity = self._client._entity_cache[self._peer] - except KeyError: + self._input_entity = self._client._mb_entity_cache.get( + get_peer_id(self._peer, add_mark=False))._as_input_peer() + except AttributeError: pass return self._input_entity diff --git a/telethon/tl/custom/forward.py b/telethon/tl/custom/forward.py index a95eae30..afe6b3a0 100644 --- a/telethon/tl/custom/forward.py +++ b/telethon/tl/custom/forward.py @@ -36,12 +36,12 @@ class Forward(ChatGetter, SenderGetter): if ty == helpers._EntityType.USER: sender_id = utils.get_peer_id(original.from_id) sender, input_sender = utils._get_entity_pair( - sender_id, entities, client._entity_cache) + sender_id, entities, client._mb_entity_cache) elif ty in (helpers._EntityType.CHAT, helpers._EntityType.CHANNEL): peer = original.from_id chat, input_chat = utils._get_entity_pair( - utils.get_peer_id(peer), entities, client._entity_cache) + utils.get_peer_id(peer), entities, client._mb_entity_cache) # This call resets the client ChatGetter.__init__(self, peer, chat=chat, input_chat=input_chat) diff --git a/telethon/tl/custom/message.py b/telethon/tl/custom/message.py index 3f20c61e..4959ed0f 100644 --- a/telethon/tl/custom/message.py +++ b/telethon/tl/custom/message.py @@ -285,7 +285,7 @@ class Message(ChatGetter, SenderGetter, TLObject): if self.peer_id == types.PeerUser(client._self_id) and not self.fwd_from: self.out = True - cache = client._entity_cache + cache = client._mb_entity_cache self._sender, self._input_sender = utils._get_entity_pair( self.sender_id, entities, cache) @@ -1138,8 +1138,9 @@ class Message(ChatGetter, SenderGetter, TLObject): return bot else: try: - return self._client._entity_cache[self.via_bot_id] - except KeyError: + return self._client._mb_entity_cache.get( + utils.resolve_id(self.via_bot_id)[0])._as_input_peer() + except AttributeError: raise ValueError('No input sender') from None def _document_by_attribute(self, kind, condition=None): diff --git a/telethon/tl/custom/sendergetter.py b/telethon/tl/custom/sendergetter.py index aa1664d3..7d389ec7 100644 --- a/telethon/tl/custom/sendergetter.py +++ b/telethon/tl/custom/sendergetter.py @@ -1,5 +1,7 @@ import abc +from ... import utils + class SenderGetter(abc.ABC): """ @@ -69,9 +71,9 @@ class SenderGetter(abc.ABC): """ if self._input_sender is None and self._sender_id and self._client: try: - self._input_sender = \ - self._client._entity_cache[self._sender_id] - except KeyError: + self._input_sender = self._client._mb_entity_cache.get( + utils.resolve_id(self._sender_id)[0])._as_input_peer() + except AttributeError: pass return self._input_sender diff --git a/telethon/utils.py b/telethon/utils.py index 0663e58c..a033d3ef 100644 --- a/telethon/utils.py +++ b/telethon/utils.py @@ -585,9 +585,9 @@ def _get_entity_pair(entity_id, entities, cache, """ entity = entities.get(entity_id) try: - input_entity = cache[entity_id] - except KeyError: - # KeyError is unlikely, so another TypeError won't hurt + input_entity = cache.get(resolve_id(entity_id)[0])._as_input_peer() + except AttributeError: + # AttributeError is unlikely, so another TypeError won't hurt try: input_entity = get_input_peer(entity) except TypeError: