Make send_media_group return the new Messages object

This commit is contained in:
Dan 2019-02-05 11:50:32 +01:00
parent 6109129f73
commit a2263ad8ce

View File

@ -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,14 +74,21 @@ class SendMediaGroup(BaseClient):
if isinstance(i, pyrogram.InputMediaPhoto):
if os.path.exists(i.media):
media = self.send(
functions.messages.UploadMedia(
peer=self.resolve_peer(chat_id),
media=types.InputMediaUploadedPhoto(
file=self.save_file(i.media)
while True:
try:
media = self.send(
functions.messages.UploadMedia(
peer=self.resolve_peer(chat_id),
media=types.InputMediaUploadedPhoto(
file=self.save_file(i.media)
)
)
)
)
)
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,25 +122,32 @@ class SendMediaGroup(BaseClient):
)
elif isinstance(i, pyrogram.InputMediaVideo):
if os.path.exists(i.media):
media = self.send(
functions.messages.UploadMedia(
peer=self.resolve_peer(chat_id),
media=types.InputMediaUploadedDocument(
file=self.save_file(i.media),
thumb=None if i.thumb is None else self.save_file(i.thumb),
mime_type=mimetypes.types_map[".mp4"],
attributes=[
types.DocumentAttributeVideo(
supports_streaming=i.supports_streaming or None,
duration=i.duration,
w=i.width,
h=i.height
),
types.DocumentAttributeFilename(os.path.basename(i.media))
]
while True:
try:
media = self.send(
functions.messages.UploadMedia(
peer=self.resolve_peer(chat_id),
media=types.InputMediaUploadedDocument(
file=self.save_file(i.media),
thumb=None if i.thumb is None else self.save_file(i.thumb),
mime_type=mimetypes.types_map[".mp4"],
attributes=[
types.DocumentAttributeVideo(
supports_streaming=i.supports_streaming or None,
duration=i.duration,
w=i.width,
h=i.height
),
types.DocumentAttributeFilename(os.path.basename(i.media))
]
)
)
)
)
)
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,11 +188,30 @@ class SendMediaGroup(BaseClient):
)
)
return self.send(
functions.messages.SendMultiMedia(
peer=self.resolve_peer(chat_id),
multi_media=multi_media,
silent=disable_notification or None,
reply_to_msg_id=reply_to_message_id
while True:
try:
r = self.send(
functions.messages.SendMultiMedia(
peer=self.resolve_peer(chat_id),
multi_media=multi_media,
silent=disable_notification or None,
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
)
)