Make send_media_group return the new Messages object

pyrogram/client/methods/messages/send_media_group.py

@@ -17,20 +17,23 @@
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 import binascii
+import logging
 import mimetypes
 import os
 import struct
+import time
 from typing import Union, List
 
 import pyrogram
 from pyrogram.api import functions, types
-from pyrogram.api.errors import FileIdInvalid
+from pyrogram.api.errors import FileIdInvalid, FloodWait
 from pyrogram.client.ext import BaseClient, utils
 
+log = logging.getLogger(__name__)
+
 
 class SendMediaGroup(BaseClient):
     # TODO: Add progress parameter
-    # TODO: Return new Message object
     # TODO: Figure out how to send albums using URLs
     def send_media_group(self,
                          chat_id: Union[int, str],
@@ -38,7 +41,6 @@ class SendMediaGroup(BaseClient):
                          disable_notification: bool = None,
                          reply_to_message_id: int = None):
         """Use this method to send a group of photos or videos as an album.
-        On success, an Update containing the sent Messages is returned.
 
         Args:
             chat_id (``int`` | ``str``):
@@ -57,6 +59,13 @@ class SendMediaGroup(BaseClient):
 
             reply_to_message_id (``int``, *optional*):
                 If the message is a reply, ID of the original message.
+
+        Returns:
+            On success, a :obj:`Messages <pyrogram.Messages>` object is returned containing all the
+            single messages sent.
+
+        Raises:
+            :class:`Error <pyrogram.Error>` in case of a Telegram RPC error.
         """
         multi_media = []
 
@@ -65,6 +74,8 @@ class SendMediaGroup(BaseClient):
 
             if isinstance(i, pyrogram.InputMediaPhoto):
                 if os.path.exists(i.media):
+                    while True:
+                        try:
                             media = self.send(
                                 functions.messages.UploadMedia(
                                     peer=self.resolve_peer(chat_id),
@@ -73,6 +84,11 @@ class SendMediaGroup(BaseClient):
                                     )
                                 )
                             )
+                        except FloodWait as e:
+                            log.warning("Sleeping for {}s".format(e.x))
+                            time.sleep(e.x)
+                        else:
+                            break
 
                     media = types.InputMediaPhoto(
                         id=types.InputPhoto(
@@ -106,6 +122,8 @@ class SendMediaGroup(BaseClient):
                         )
             elif isinstance(i, pyrogram.InputMediaVideo):
                 if os.path.exists(i.media):
+                    while True:
+                        try:
                             media = self.send(
                                 functions.messages.UploadMedia(
                                     peer=self.resolve_peer(chat_id),
@@ -125,6 +143,11 @@ class SendMediaGroup(BaseClient):
                                     )
                                 )
                             )
+                        except FloodWait as e:
+                            log.warning("Sleeping for {}s".format(e.x))
+                            time.sleep(e.x)
+                        else:
+                            break
 
                     media = types.InputMediaDocument(
                         id=types.InputDocument(
@@ -165,7 +188,9 @@ class SendMediaGroup(BaseClient):
                 )
             )
 
-        return self.send(
+        while True:
+            try:
+                r = self.send(
                     functions.messages.SendMultiMedia(
                         peer=self.resolve_peer(chat_id),
                         multi_media=multi_media,
@@ -173,3 +198,20 @@ class SendMediaGroup(BaseClient):
                         reply_to_msg_id=reply_to_message_id
                     )
                 )
+            except FloodWait as e:
+                log.warning("Sleeping for {}s".format(e.x))
+                time.sleep(e.x)
+            else:
+                break
+
+        return pyrogram.Messages._parse(
+            self,
+            types.messages.Messages(
+                messages=[m.message for m in filter(
+                    lambda u: isinstance(u, (types.UpdateNewMessage, types.UpdateNewChannelMessage)),
+                    r.updates
+                )],
+                users=r.users,
+                chats=r.chats
+            )
+        )