Update forward_messages to use _get_response_message

telethon/client/messageparse.py

@@ -93,19 +93,9 @@ class MessageParseMethods(UserMethods):
         """
         Extracts the response message known a request and Update result.
         The request may also be the ID of the message to match.
-        """
-        # Telegram seems to send updateMessageID first, then updateNewMessage,
-        # however let's not rely on that just in case.
-        if isinstance(request, int):
-            msg_id = request
-        else:
-            msg_id = None
-            for update in result.updates:
-                if isinstance(update, types.UpdateMessageID):
-                    if update.random_id == request.random_id:
-                        msg_id = update.id
-                        break
 
+        If ``request.random_id`` is a list, this method returns a list too.
+        """
         if isinstance(result, types.UpdateShort):
             updates = [result.update]
             entities = {}
@@ -117,31 +107,43 @@ class MessageParseMethods(UserMethods):
         else:
             return None
 
-        found = None
+        random_to_id = {}
+        id_to_message = {}
         for update in updates:
-            if isinstance(update, (
+            if isinstance(update, types.UpdateMessageID):
+                random_to_id[update.random_id] = update.id
+
+            elif isinstance(update, (
                     types.UpdateNewChannelMessage, types.UpdateNewMessage)):
-                if update.message.id == msg_id:
+                update.message._finish_init(self, entities, input_chat)
-                    found = update.message
+                id_to_message[update.message.id] = update.message
-                    break
 
             elif (isinstance(update, types.UpdateEditMessage)
                   and not isinstance(request.peer, types.InputPeerChannel)):
                 if request.id == update.message.id:
-                    found = update.message
+                    update.message._finish_init(self, entities, input_chat)
-                    break
+                    return update.message
 
             elif (isinstance(update, types.UpdateEditChannelMessage)
                   and utils.get_peer_id(request.peer) ==
-                    utils.get_peer_id(update.message.to_id)):
+                  utils.get_peer_id(update.message.to_id)):
                 if request.id == update.message.id:
-                    found = update.message
+                    update.message._finish_init(self, entities, input_chat)
-                    break
+                    return update.message
 
-        if found:
+        if not utils.is_list_like(request.random_id):
-            found._finish_init(self, entities, input_chat)
+            if request.random_id in random_to_id:
-            return found
+                return id_to_message[random_to_id[request.random_id]]
+            else:
+                return None
         else:
-            return None  # explicit is better than implicit
+            # ``rnd in random_to_id`` is needed because trying to forward only
+            # deleted messages causes `MESSAGE_ID_INVALID`, but forwarding
+            # valid and invalid messages in the same call makes the call
+            # succeed, although the API won't return those messages thus
+            # `random_to_id[rnd]` would `KeyError`.
+            return [id_to_message[random_to_id[rnd]]
+                    if rnd in random_to_id else None
+                    for rnd in request.random_id]
 
     # endregion

telethon/client/messages.py

@@ -667,31 +667,8 @@ class MessageMethods(UploadMethods, ButtonMethods, MessageParseMethods):
             silent=silent
         )
         result = await self(req)
-        if isinstance(result, (types.Updates, types.UpdatesCombined)):
+        sent = self._get_response_message(req, result, entity)
-            entities = {utils.get_peer_id(x): x
+        return sent[0] if single else sent
-                        for x in itertools.chain(result.users, result.chats)}
-        else:
-            entities = {}
-
-        random_to_id = {}
-        id_to_message = {}
-        for update in result.updates:
-            if isinstance(update, types.UpdateMessageID):
-                random_to_id[update.random_id] = update.id
-            elif isinstance(update, (
-                    types.UpdateNewMessage, types.UpdateNewChannelMessage)):
-                update.message._finish_init(self, entities, entity)
-                id_to_message[update.message.id] = update.message
-
-        # Trying to forward only deleted messages causes `MESSAGE_ID_INVALID`
-        # but forwarding valid and invalid messages in the same call makes the
-        # call succeed, although the API won't return those messages thus
-        # `random_to_id[rnd]` would `KeyError`. Check the key beforehand.
-        result = [id_to_message[random_to_id[rnd]]
-                  if rnd in random_to_id else None
-                  for rnd in req.random_id]
-
-        return result[0] if single else result
 
     async def edit_message(
             self, entity, message=None, text=None,