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",
"new_chat_photo")
if isinstance(message, (types.Message, types.Story)):
story = getattr(message, "story", None)
media = None
if isinstance(message, types.Message) and message.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:
break
else:
raise ValueError("This message doesn't contain any downloadable media")
else:
elif isinstance(message, types.Story):
media = getattr(message, message.media.value, None)
elif isinstance(message, str):
media = message
if not media:
raise ValueError("This message doesn't contain any downloadable media")
if isinstance(media, str):
file_id_str = media
else:

View File

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