diff --git a/docs/source/releases/changes-in-this-fork.rst b/docs/source/releases/changes-in-this-fork.rst index 9cf18d556..d05650f1f 100644 --- a/docs/source/releases/changes-in-this-fork.rst +++ b/docs/source/releases/changes-in-this-fork.rst @@ -14,6 +14,7 @@ If you found any issue or have any suggestions, feel free to make `an issue `_ - PR `#115 `_ This `change `_ breaks some usages with offset-naive and offset-aware datetimes. - PR from upstream: `#1411 `_ without attribution. diff --git a/pyrogram/methods/messages/download_media.py b/pyrogram/methods/messages/download_media.py index 31a5a7f78..b634da1a4 100644 --- a/pyrogram/methods/messages/download_media.py +++ b/pyrogram/methods/messages/download_media.py @@ -19,11 +19,12 @@ import asyncio import io import os +import re from datetime import datetime -from typing import Union, Optional, Callable +from typing import Callable, Optional, Union import pyrogram -from pyrogram import types, enums +from pyrogram import enums, types, utils from pyrogram.file_id import FileId, FileType, PHOTO_TYPES DEFAULT_DOWNLOAD_DIR = "downloads/" @@ -32,19 +33,39 @@ class DownloadMedia: async def download_media( self: "pyrogram.Client", - message: Union["types.Message", "types.Story", str], + message: Union[ + "types.Message", + "types.Audio", + "types.Document", + "types.Photo", + "types.Sticker", + "types.Video", + "types.Animation", + "types.Voice", + "types.VideoNote", + # TODO + "types.Story", + "types.PaidMediaInfo", + "types.PaidMediaPhoto", + "types.PaidMediaVideo", + "types.Thumbnail", + "types.StrippedThumbnail", + "types.PaidMediaPreview", + str, + ], file_name: str = DEFAULT_DOWNLOAD_DIR, in_memory: bool = False, block: bool = True, + idx: int = None, progress: Callable = None, progress_args: tuple = () - ) -> Optional[Union[str, "io.BytesIO"]]: + ) -> Optional[Union[str, "io.BytesIO", list[str], list["io.BytesIO"]]]: """Download the media from a message. .. include:: /_includes/usable-by/users-bots.rst Parameters: - message (:obj:`~pyrogram.types.Message` | :obj:`~pyrogram.types.Story` | ``str``): + message (:obj:`~pyrogram.types.Message` | :obj:`~pyrogram.types.Audio` | :obj:`~pyrogram.types.Document` | :obj:`~pyrogram.types.Photo` | :obj:`~pyrogram.types.Sticker` | :obj:`~pyrogram.types.Video` | :obj:`~pyrogram.types.Animation` | :obj:`~pyrogram.types.Voice` | :obj:`~pyrogram.types.VideoNote` | :obj:`~pyrogram.types.Story` | :obj:`~pyrogram.types.PaidMediaInfo` | :obj:`~pyrogram.types.PaidMediaPhoto` | :obj:`~pyrogram.types.PaidMediaVideo` | :obj:`~pyrogram.types.Thumbnail` | :obj:`~pyrogram.types.StrippedThumbnail` | :obj:`~pyrogram.types.PaidMediaPreview` | :obj:`~pyrogram.types.Story` | ``str``): Pass a Message containing the media, the media itself (message.audio, message.video, ...) or a file id as string. @@ -63,6 +84,9 @@ async def download_media( Blocks the code execution until the file has been downloaded. Defaults to True. + idx (``int``, *optional*): + In case of a :obj:`~pyrogram.types.PaidMediaInfo` with more than one ``paid_media``, the zero based index of the :obj:`~pyrogram.types.PaidMedia` to download. Raises ``IndexError`` if the index specified does not exist in the original ``message``. + progress (``Callable``, *optional*): Pass a callback function to view the file transmission progress. The function must take *(current, total)* as positional arguments (look at Other Parameters below for a @@ -90,9 +114,11 @@ async def download_media( otherwise, in case the download failed or was deliberately stopped with :meth:`~pyrogram.Client.stop_transmission`, None is returned. Otherwise, in case ``in_memory=True``, a binary file-like object with its attribute ".name" set is returned. + If the message is a :obj:`~pyrogram.types.PaidMediaInfo` with more than one ``paid_media`` containing ``minithumbnail`` and ``idx`` is not specified, then a list of paths or binary file-like objects is returned. Raises: RPCError: In case of a Telegram RPC error. + IndexError: In case of wrong value of ``idx``. ValueError: If the message doesn't contain any downloadable media. Example: @@ -122,11 +148,11 @@ async def progress(current, total): file_bytes = bytes(file.getbuffer()) """ - media = message + medium = [message] if isinstance(message, types.Message): if message.new_chat_photo: - media = message.new_chat_photo + medium = [message.new_chat_photo] elif ( not (self.me and self.me.is_bot) and @@ -134,82 +160,149 @@ async def progress(current, total): ): story_media = message.story or message.reply_to_story or None if story_media and story_media.media: - media = getattr(story_media, story_media.media.value, None) + medium = [getattr(story_media, story_media.media.value, None)] + else: + medium = [] + + elif message.paid_media: + if any([isinstance(paid_media, (types.PaidMediaPhoto, types.PaidMediaVideo)) for paid_media in message.paid_media.paid_media]): + medium = [getattr(paid_media, "photo", (getattr(paid_media, "video", None))) for paid_media in message.paid_media.paid_media] + elif any([isinstance(paid_media, types.PaidMediaPreview) for paid_media in message.paid_media.paid_media]): + medium = [getattr(getattr(paid_media, "minithumbnail"), "data", None) for paid_media in message.paid_media.paid_media] else: - media = None + medium = [] else: if message.media: - media = getattr(message, message.media.value, None) + medium = [getattr(message, message.media.value, None)] else: - media = None + medium = [] - elif isinstance(message, str): - media = message - - if isinstance(media, types.Story): + elif isinstance(message, types.Story): if (self.me and self.me.is_bot): raise ValueError("This method cannot be used by bots") else: - if media.media: - media = getattr(message, message.media.value, None) + if medium.media: + medium = [getattr(message, message.media.value, None)] else: - media = None + medium = [] + + elif isinstance(message, types.PaidMediaInfo): + if any([isinstance(paid_media, (types.PaidMediaPhoto, types.PaidMediaVideo)) for paid_media in message.paid_media]): + medium = [getattr(paid_media, "photo", (getattr(paid_media, "video", None))) for paid_media in message.paid_media] + elif any([isinstance(paid_media, types.PaidMediaPreview) for paid_media in message.paid_media]): + medium = [getattr(getattr(paid_media, "minithumbnail"), "data", None) for paid_media in message.paid_media] + else: + medium = [] + + elif isinstance(message, types.PaidMediaPhoto): + medium = [message.photo] + + elif isinstance(message, types.PaidMediaVideo): + medium = [message.video] + + elif isinstance(message, types.PaidMediaPreview): + medium = [getattr(getattr(message, "minithumbnail"), "data", None)] + + elif isinstance(message, types.StrippedThumbnail): + medium = [message.data] + + elif isinstance(message, types.Thumbnail): + medium = [message] + + elif isinstance(message, str): + medium = [message] - if not media: + medium = types.List(filter(lambda x: x is not None, medium)) + + if len(medium) == 0: raise ValueError( f"The message {message if isinstance(message, str) else message.id} doesn't contain any downloadable media" ) - if isinstance(media, str): - file_id_str = media - else: - file_id_str = media.file_id - - file_id_obj = FileId.decode(file_id_str) - - file_type = file_id_obj.file_type - media_file_name = getattr(media, "file_name", "") # TODO - file_size = getattr(media, "file_size", 0) - mime_type = getattr(media, "mime_type", "") - date = getattr(media, "date", None) - - directory, file_name = os.path.split(file_name) - file_name = file_name or media_file_name or "" - - if not os.path.isabs(file_name): - directory = self.WORKDIR / (directory or DEFAULT_DOWNLOAD_DIR) - - if not file_name: - guessed_extension = self.guess_extension(mime_type) - - if file_type in PHOTO_TYPES: - extension = ".jpg" - elif file_type == FileType.VOICE: - extension = guessed_extension or ".ogg" - elif file_type in (FileType.VIDEO, FileType.ANIMATION, FileType.VIDEO_NOTE): - extension = guessed_extension or ".mp4" - elif file_type == FileType.DOCUMENT: - extension = guessed_extension or ".zip" - elif file_type == FileType.STICKER: - extension = guessed_extension or ".webp" - elif file_type == FileType.AUDIO: - extension = guessed_extension or ".mp3" + if idx is not None: + medium = [medium[idx]] + + dledmedia = [] + + for media in medium: + if isinstance(media, bytes): + thumb = utils.from_inline_bytes( + utils.expand_inline_bytes( + media + ) + ) + if in_memory: + dledmedia.append(thumb) + continue + + directory, file_name = os.path.split(file_name) + file_name = file_name or thumb.name + + if not os.path.isabs(file_name): + directory = self.PARENT_DIR / (directory or DEFAULT_DOWNLOAD_DIR) + + os.makedirs(directory, exist_ok=True) if not in_memory else None + temp_file_path = os.path.abspath(re.sub("\\\\", "/", os.path.join(directory, file_name))) + + with open(temp_file_path, "wb") as file: + file.write(thumb.getbuffer()) + + dledmedia.append(temp_file_path) + continue + + elif isinstance(media, str): + file_id_str = media else: - extension = ".unknown" + file_id_str = media.file_id + + file_id_obj = FileId.decode(file_id_str) + + file_type = file_id_obj.file_type + media_file_name = getattr(media, "file_name", "") # TODO + file_size = getattr(media, "file_size", 0) + mime_type = getattr(media, "mime_type", "") + date = getattr(media, "date", None) + + directory, file_name = os.path.split(file_name) + # TODO + file_name = file_name or media_file_name or "" + + if not os.path.isabs(file_name): + directory = self.WORKDIR / (directory or DEFAULT_DOWNLOAD_DIR) + + if not file_name: + guessed_extension = self.guess_extension(mime_type) + + if file_type in PHOTO_TYPES: + extension = ".jpg" + elif file_type == FileType.VOICE: + extension = guessed_extension or ".ogg" + elif file_type in (FileType.VIDEO, FileType.ANIMATION, FileType.VIDEO_NOTE): + extension = guessed_extension or ".mp4" + elif file_type == FileType.DOCUMENT: + extension = guessed_extension or ".zip" + elif file_type == FileType.STICKER: + extension = guessed_extension or ".webp" + elif file_type == FileType.AUDIO: + extension = guessed_extension or ".mp3" + else: + extension = ".unknown" + + file_name = "{}_{}_{}{}".format( + FileType(file_id_obj.file_type).name.lower(), + (date or datetime.now()).strftime("%Y-%m-%d_%H-%M-%S"), + self.rnd_id(), + extension + ) - file_name = "{}_{}_{}{}".format( - FileType(file_id_obj.file_type).name.lower(), - (date or datetime.now()).strftime("%Y-%m-%d_%H-%M-%S"), - self.rnd_id(), - extension + downloader = self.handle_download( + (file_id_obj, directory, file_name, in_memory, file_size, progress, progress_args) ) - downloader = self.handle_download( - (file_id_obj, directory, file_name, in_memory, file_size, progress, progress_args) - ) + if block: + dledmedia.append(await downloader) + else: + asyncio.get_event_loop().create_task(downloader) - if block: - return await downloader - else: - asyncio.get_event_loop().create_task(downloader) + return types.List(dledmedia) if block and len(dledmedia) > 1 else dledmedia[0] if block and len(dledmedia) == 1 else None diff --git a/pyrogram/types/input_paid_media/paid_media.py b/pyrogram/types/input_paid_media/paid_media.py index f4e5c3b7b..6e532e25c 100644 --- a/pyrogram/types/input_paid_media/paid_media.py +++ b/pyrogram/types/input_paid_media/paid_media.py @@ -55,7 +55,8 @@ def _parse( duration=getattr(extended_media, "video_duration", None), minithumbnail=types.StrippedThumbnail( client=client, - data=extended_media.thumb + # TODO + data=getattr(getattr(extended_media, "thumb"), "bytes", None) ) if getattr(extended_media, "thumb", None) else None ) if isinstance(extended_media, raw.types.MessageExtendedMedia): diff --git a/pyrogram/types/messages_and_media/message.py b/pyrogram/types/messages_and_media/message.py index 2030ff98f..09260c9dc 100644 --- a/pyrogram/types/messages_and_media/message.py +++ b/pyrogram/types/messages_and_media/message.py @@ -5298,9 +5298,10 @@ async def download( file_name: str = "", in_memory: bool = False, block: bool = True, + idx: int = None, progress: Callable = None, progress_args: tuple = () - ) -> Union[str, "io.BytesIO"]: + ) -> Optional[Union[str, "io.BytesIO", list[str], list["io.BytesIO"]]]: """Bound method *download* of :obj:`~pyrogram.types.Message`. Use as a shortcut for: @@ -5330,6 +5331,9 @@ async def download( Blocks the code execution until the file has been downloaded. Defaults to True. + idx (``int``, *optional*): + In case of a :obj:`~pyrogram.types.PaidMediaInfo` with more than one ``paid_media``, the zero based index of the :obj:`~pyrogram.types.PaidMedia` to download. Raises ``IndexError`` if the index specified does not exist in the original ``message``. + progress (``Callable``, *optional*): Pass a callback function to view the file transmission progress. The function must take *(current, total)* as positional arguments (look at Other Parameters below for a @@ -5357,9 +5361,11 @@ async def download( otherwise, in case the download failed or was deliberately stopped with :meth:`~pyrogram.Client.stop_transmission`, None is returned. Otherwise, in case ``in_memory=True``, a binary file-like object with its attribute ".name" set is returned. + If the message is a :obj:`~pyrogram.types.PaidMediaInfo` with more than one ``paid_media`` containing ``minithumbnail`` and ``idx`` is not specified, then a list of paths or binary file-like objects is returned. Raises: RPCError: In case of a Telegram RPC error. + IndexError: In case of wrong value of ``idx``. ValueError: If the message doesn't contain any downloadable media. """ @@ -5368,6 +5374,7 @@ async def download( file_name=file_name, in_memory=in_memory, block=block, + idx=idx, progress=progress, progress_args=progress_args, ) diff --git a/pyrogram/types/messages_and_media/story.py b/pyrogram/types/messages_and_media/story.py index 29853cf9a..adb712dc6 100644 --- a/pyrogram/types/messages_and_media/story.py +++ b/pyrogram/types/messages_and_media/story.py @@ -16,8 +16,9 @@ # You should have received a copy of the GNU Lesser General Public License # along with Pyrogram. If not, see . +import io from datetime import datetime -from typing import Union, Callable +from typing import Callable, Optional, Union import pyrogram from pyrogram import raw, utils, types, enums @@ -453,7 +454,7 @@ async def download( block: bool = True, progress: Callable = None, progress_args: tuple = () - ) -> str: + ) -> Optional[Union[str, "io.BytesIO"]]: """Bound method *download* of :obj:`~pyrogram.types.Story`. Use as a shortcut for: @@ -506,7 +507,10 @@ async def download( You can either keep ``*args`` or add every single extra argument in your function signature. Returns: - On success, the absolute path of the downloaded file as string is returned, None otherwise. + ``str`` | ``None`` | :obj:`io.BytesIO`: On success, the absolute path of the downloaded file is returned, + otherwise, in case the download failed or was deliberately stopped with + :meth:`~pyrogram.Client.stop_transmission`, None is returned. + Otherwise, in case ``in_memory=True``, a binary file-like object with its attribute ".name" set is returned. Raises: RPCError: In case of a Telegram RPC error. diff --git a/pyrogram/utils.py b/pyrogram/utils.py index 759383671..a429deac0 100644 --- a/pyrogram/utils.py +++ b/pyrogram/utils.py @@ -26,6 +26,7 @@ from concurrent.futures.thread import ThreadPoolExecutor from datetime import datetime, timezone from getpass import getpass +from io import BytesIO from typing import Union, Optional import pyrogram @@ -592,3 +593,58 @@ def fix_up_voice_audio_uri( mime_type = "audio/ogg" # BEWARE: https://t.me/c/1279877202/74 return mime_type + + +def expand_inline_bytes(bytes_data: bytes): + # https://github.com/telegramdesktop/tdesktop/blob/1757dd856/Telegram/SourceFiles/ui/image/image.cpp#L43-L94 + if len(bytes_data) < 3 or bytes_data[0] != 0x01: + return bytearray() + header = bytearray( + b"\xff\xd8\xff\xe0\x00\x10\x4a\x46\x49" + b"\x46\x00\x01\x01\x00\x00\x01\x00\x01\x00\x00\xff\xdb\x00\x43\x00\x28\x1c" + b"\x1e\x23\x1e\x19\x28\x23\x21\x23\x2d\x2b\x28\x30\x3c\x64\x41\x3c\x37\x37" + b"\x3c\x7b\x58\x5d\x49\x64\x91\x80\x99\x96\x8f\x80\x8c\x8a\xa0\xb4\xe6\xc3" + b"\xa0\xaa\xda\xad\x8a\x8c\xc8\xff\xcb\xda\xee\xf5\xff\xff\xff\x9b\xc1\xff" + b"\xff\xff\xfa\xff\xe6\xfd\xff\xf8\xff\xdb\x00\x43\x01\x2b\x2d\x2d\x3c\x35" + b"\x3c\x76\x41\x41\x76\xf8\xa5\x8c\xa5\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8" + b"\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8" + b"\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8\xf8" + b"\xf8\xf8\xf8\xf8\xf8\xff\xc0\x00\x11\x08\x00\x00\x00\x00\x03\x01\x22\x00" + b"\x02\x11\x01\x03\x11\x01\xff\xc4\x00\x1f\x00\x00\x01\x05\x01\x01\x01\x01" + b"\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x01\x02\x03\x04\x05\x06\x07\x08" + b"\x09\x0a\x0b\xff\xc4\x00\xb5\x10\x00\x02\x01\x03\x03\x02\x04\x03\x05\x05" + b"\x04\x04\x00\x00\x01\x7d\x01\x02\x03\x00\x04\x11\x05\x12\x21\x31\x41\x06" + b"\x13\x51\x61\x07\x22\x71\x14\x32\x81\x91\xa1\x08\x23\x42\xb1\xc1\x15\x52" + b"\xd1\xf0\x24\x33\x62\x72\x82\x09\x0a\x16\x17\x18\x19\x1a\x25\x26\x27\x28" + b"\x29\x2a\x34\x35\x36\x37\x38\x39\x3a\x43\x44\x45\x46\x47\x48\x49\x4a\x53" + b"\x54\x55\x56\x57\x58\x59\x5a\x63\x64\x65\x66\x67\x68\x69\x6a\x73\x74\x75" + b"\x76\x77\x78\x79\x7a\x83\x84\x85\x86\x87\x88\x89\x8a\x92\x93\x94\x95\x96" + b"\x97\x98\x99\x9a\xa2\xa3\xa4\xa5\xa6\xa7\xa8\xa9\xaa\xb2\xb3\xb4\xb5\xb6" + b"\xb7\xb8\xb9\xba\xc2\xc3\xc4\xc5\xc6\xc7\xc8\xc9\xca\xd2\xd3\xd4\xd5\xd6" + b"\xd7\xd8\xd9\xda\xe1\xe2\xe3\xe4\xe5\xe6\xe7\xe8\xe9\xea\xf1\xf2\xf3\xf4" + b"\xf5\xf6\xf7\xf8\xf9\xfa\xff\xc4\x00\x1f\x01\x00\x03\x01\x01\x01\x01\x01" + b"\x01\x01\x01\x01\x00\x00\x00\x00\x00\x00\x01\x02\x03\x04\x05\x06\x07\x08" + b"\x09\x0a\x0b\xff\xc4\x00\xb5\x11\x00\x02\x01\x02\x04\x04\x03\x04\x07\x05" + b"\x04\x04\x00\x01\x02\x77\x00\x01\x02\x03\x11\x04\x05\x21\x31\x06\x12\x41" + b"\x51\x07\x61\x71\x13\x22\x32\x81\x08\x14\x42\x91\xa1\xb1\xc1\x09\x23\x33" + b"\x52\xf0\x15\x62\x72\xd1\x0a\x16\x24\x34\xe1\x25\xf1\x17\x18\x19\x1a\x26" + b"\x27\x28\x29\x2a\x35\x36\x37\x38\x39\x3a\x43\x44\x45\x46\x47\x48\x49\x4a" + b"\x53\x54\x55\x56\x57\x58\x59\x5a\x63\x64\x65\x66\x67\x68\x69\x6a\x73\x74" + b"\x75\x76\x77\x78\x79\x7a\x82\x83\x84\x85\x86\x87\x88\x89\x8a\x92\x93\x94" + b"\x95\x96\x97\x98\x99\x9a\xa2\xa3\xa4\xa5\xa6\xa7\xa8\xa9\xaa\xb2\xb3\xb4" + b"\xb5\xb6\xb7\xb8\xb9\xba\xc2\xc3\xc4\xc5\xc6\xc7\xc8\xc9\xca\xd2\xd3\xd4" + b"\xd5\xd6\xd7\xd8\xd9\xda\xe2\xe3\xe4\xe5\xe6\xe7\xe8\xe9\xea\xf2\xf3\xf4" + b"\xf5\xf6\xf7\xf8\xf9\xfa\xff\xda\x00\x0c\x03\x01\x00\x02\x11\x03\x11\x00" + b"\x3f\x00" + ) + footer = bytearray(b"\xff\xd9") + header[164] = bytes_data[1] + header[166] = bytes_data[2] + return header + bytes_data[3:] + footer + + +def from_inline_bytes(data: bytes, file_name: str = None) -> BytesIO: + b = BytesIO() + b.write(data) + b.name = file_name if file_name else f"photo_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.jpg" + return b