Reuse code for _get_entity_pair

Less error-prone, improved the function flow for all callers,
and removed some duplicate work.

telethon/events/callbackquery.py

@@ -105,7 +105,8 @@ def __init__(self, query, peer, msg_id):
 
         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):

telethon/events/common.py

@@ -147,25 +147,11 @@ def _set_client(self, client):
         """
         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):
         """

telethon/events/inlinequery.py

@@ -89,7 +89,8 @@ def __init__(self, query):
 
         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):

telethon/events/newmessage.py

@@ -206,16 +206,6 @@ def _set_client(self, client):
             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):

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 @@ def _set_client(self, client):
                     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):

telethon/tl/custom/forward.py

@@ -29,25 +29,16 @@ def __init__(self, client, original, entities):
         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))
+
+            self._chat, self._input_chat = utils._get_entity_pair(
+                 self.chat_id, entities, client._entity_cache)
         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 = self._chat_peer = None
 
     # TODO We could reload the message

telethon/tl/custom/message.py

@@ -210,27 +210,20 @@ def _finish_init(self, client, entities, input_chat):
         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)

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:

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'