Refactor download_media

This commit is contained in:
KurimuzonAkuma 2024-03-10 23:14:00 +03:00
parent 6f99be7ae6
commit 8df8780841
2 changed files with 14 additions and 7 deletions

View File

@ -122,19 +122,26 @@ class DownloadMedia:
available_media = ("audio", "document", "photo", "sticker", "animation", "video", "voice", "video_note", available_media = ("audio", "document", "photo", "sticker", "animation", "video", "voice", "video_note",
"new_chat_photo") "new_chat_photo")
if isinstance(message, (types.Message, types.Story)): media = None
story = getattr(message, "story", None)
if isinstance(message, types.Message) and message.media:
for kind in available_media: for kind in available_media:
media = getattr(story or message, kind, None) story = message.story or message.reply_to_story
if story:
media = getattr(story, kind, None)
else:
media = getattr(message, kind, None)
if media is not None: if media is not None:
break break
else: elif isinstance(message, types.Story):
raise ValueError("This message doesn't contain any downloadable media") media = getattr(message, message.media.value, None)
else: elif isinstance(message, str):
media = message media = message
if not media:
raise ValueError("This message doesn't contain any downloadable media")
if isinstance(media, str): if isinstance(media, str):
file_id_str = media file_id_str = media
else: else:

View File

@ -1788,7 +1788,7 @@ class Story(Object, Update):
``ValueError``: If the message doesn't contain any downloadable media ``ValueError``: If the message doesn't contain any downloadable media
""" """
return await self._client.download_media( return await self._client.download_media(
message=getattr(self, self.media.value), message=self,
file_name=file_name, file_name=file_name,
in_memory=in_memory, in_memory=in_memory,
block=block, block=block,