diff --git a/telethon/events/callbackquery.py b/telethon/events/callbackquery.py index 2806f5b2..6c94cc32 100644 --- a/telethon/events/callbackquery.py +++ b/telethon/events/callbackquery.py @@ -105,7 +105,8 @@ class CallbackQuery(EventBuilder): def _set_client(self, client): super()._set_client(client) - self._sender, self._input_sender = self._get_entity_pair(self.sender_id) + self._sender, self._input_sender = utils._get_entity_pair( + self.sender_id, self._entities, client._entity_cache) @property def id(self): diff --git a/telethon/events/common.py b/telethon/events/common.py index 3e4615f9..0a4e57f8 100644 --- a/telethon/events/common.py +++ b/telethon/events/common.py @@ -147,25 +147,11 @@ class EventCommon(ChatGetter, abc.ABC): """ self._client = client if self._chat_peer: - self._chat, self._input_chat = self._get_entity_pair(self.chat_id) + self._chat, self._input_chat = utils._get_entity_pair( + self.chat_id, self._entities, client._entity_cache) else: self._chat = self._input_chat = None - def _get_entity_pair(self, entity_id): - """ - Returns ``(entity, input_entity)`` for the given entity ID. - """ - entity = self._entities.get(entity_id) - try: - input_entity = utils.get_input_peer(entity) - except TypeError: - try: - input_entity = self._client._entity_cache[entity_id] - except KeyError: - input_entity = None - - return entity, input_entity - @property def client(self): """ diff --git a/telethon/events/inlinequery.py b/telethon/events/inlinequery.py index 9eae5768..4845295f 100644 --- a/telethon/events/inlinequery.py +++ b/telethon/events/inlinequery.py @@ -89,7 +89,8 @@ class InlineQuery(EventBuilder): def _set_client(self, client): super()._set_client(client) - self._sender, self._input_sender = self._get_entity_pair(self.sender_id) + self._sender, self._input_sender = utils._get_entity_pair( + self.sender_id, self._entities, client._entity_cache) @property def id(self): diff --git a/telethon/events/newmessage.py b/telethon/events/newmessage.py index 082aeb22..1c70f7a8 100644 --- a/telethon/events/newmessage.py +++ b/telethon/events/newmessage.py @@ -206,16 +206,6 @@ class NewMessage(EventBuilder): super()._set_client(client) m = self.message m._finish_init(client, self._entities, None) - - # TODO Duplicated work here - m._chat, m._input_chat = self._get_entity_pair(m.chat_id) - m._sender, m._input_sender = self._get_entity_pair(m.sender_id) - m._via_bot, m._via_input_bot = self._get_entity_pair(m.via_bot_id) - if m.forward: - f = m.forward - f._chat, f._input_chat = self._get_entity_pair(f.chat_id) - f._sender, f._input_sender = self._get_entity_pair(f.sender_id) - self.__dict__['_init'] = True # No new attributes can be set def __getattr__(self, item): diff --git a/telethon/events/userupdate.py b/telethon/events/userupdate.py index 988feac0..63ebb4cf 100644 --- a/telethon/events/userupdate.py +++ b/telethon/events/userupdate.py @@ -1,6 +1,7 @@ import datetime from .common import EventBuilder, EventCommon, name_inner_event +from .. import utils from ..tl import types from ..tl.custom.sendergetter import SenderGetter @@ -184,7 +185,8 @@ class UserUpdate(EventBuilder): self._chat_peer = types.PeerChannel(self._chat_peer) super()._set_client(client) - self._sender, self._input_sender = self._get_entity_pair(self._sender_id) + self._sender, self._input_sender = utils._get_entity_pair( + self.sender_id, self._entities, client._entity_cache) @property def user(self): diff --git a/telethon/tl/custom/forward.py b/telethon/tl/custom/forward.py index 3a3597f7..c72e85e1 100644 --- a/telethon/tl/custom/forward.py +++ b/telethon/tl/custom/forward.py @@ -29,25 +29,16 @@ class Forward(ChatGetter, SenderGetter): self.original_fwd = original self._sender_id = original.from_id - self._sender = entities.get(original.from_id) - try: - self._input_sender =\ - utils.get_input_peer(self._sender) if self._sender else None - except TypeError: - self._input_sender = None + self._sender, self._input_sender = utils._get_entity_pair( + original.from_id, entities, client._entity_cache) self._broadcast = None if original.channel_id: self._chat_peer = types.PeerChannel(original.channel_id) - self._chat = entities.get(utils.get_peer_id(self._chat_peer)) - else: - self._chat_peer = None - self._chat = None - try: - self._input_chat = \ - utils.get_input_peer(self._chat) if self._chat else None - except TypeError: - self._input_chat = None + self._chat, self._input_chat = utils._get_entity_pair( + self.chat_id, entities, client._entity_cache) + else: + self._chat = self._input_chat = self._chat_peer = None # TODO We could reload the message diff --git a/telethon/tl/custom/message.py b/telethon/tl/custom/message.py index a6a449f8..80ab68c6 100644 --- a/telethon/tl/custom/message.py +++ b/telethon/tl/custom/message.py @@ -210,27 +210,20 @@ class Message(ChatGetter, SenderGetter, TLObject, abc.ABC): known entities. """ self._client = client - self._sender = entities.get(self._sender_id) - if self._sender: - try: - self._input_sender = utils.get_input_peer(self._sender) - except TypeError: - self._input_sender = None + cache = client._entity_cache - self._chat = entities.get(self.chat_id) - self._input_chat = input_chat - if not self._input_chat and self._chat: - try: - self._input_chat = utils.get_input_peer(self._chat) - except TypeError: - self._input_chat = None + self._sender, self._input_sender = utils._get_entity_pair( + self.sender_id, entities, cache) - self._via_bot = entities.get(self.via_bot_id) - if self._via_bot: - try: - self._via_input_bot = utils.get_input_peer(self._via_bot) - except TypeError: - self._via_input_bot = None + self._chat, self._input_chat = utils._get_entity_pair( + self.chat_id, entities, cache) + + if input_chat: # This has priority + self._input_chat = input_chat + + if self.via_bot_id: + self._via_bot, self._via_input_bot = utils._get_entity_pair( + self.via_bot_id, entities, cache) if self.fwd_from: self._forward = Forward(self._client, self.fwd_from, entities) diff --git a/telethon/utils.py b/telethon/utils.py index d6df6b85..d45ad39b 100644 --- a/telethon/utils.py +++ b/telethon/utils.py @@ -469,6 +469,24 @@ def get_input_message(message): _raise_cast_fail(message, 'InputMedia') +def _get_entity_pair(entity_id, entities, cache, + get_input_peer=get_input_peer): + """ + Returns ``(entity, input_entity)`` for the given entity ID. + """ + entity = entities.get(entity_id) + try: + input_entity = cache[entity_id] + except KeyError: + # KeyError is unlikely, so another TypeError won't hurt + try: + input_entity = get_input_peer(entity) + except TypeError: + input_entity = None + + return entity, input_entity + + def get_message_id(message): """Similar to :meth:`get_input_peer`, but for message IDs.""" if message is None: diff --git a/telethon/version.py b/telethon/version.py index 0b047156..b4ebf27f 100644 --- a/telethon/version.py +++ b/telethon/version.py @@ -1,3 +1,3 @@ # Versions should comply with PEP440. # This line is parsed in setup.py: -__version__ = '1.7.2' +__version__ = '1.7.3'