Fix .download_media() not accepting Document

This commit is contained in:
Lonami Exo 2018-01-23 12:10:23 +01:00
parent 81c95b5a60
commit 58d90e7e34
2 changed files with 16 additions and 13 deletions

View File

@ -67,7 +67,7 @@ from .tl.types import (
PeerUser, InputPeerUser, InputPeerChat, InputPeerChannel, MessageEmpty,
ChatInvite, ChatInviteAlready, PeerChannel, Photo, InputPeerSelf,
InputSingleMedia, InputMediaPhoto, InputPhoto, InputFile, InputFileBig,
InputDocument, InputMediaDocument
InputDocument, InputMediaDocument, Document
)
from .tl.types.messages import DialogsSlice
from .extensions import markdown, html
@ -1308,7 +1308,7 @@ class TelegramClient(TelegramBareClient):
return self._download_photo(
media, file, date, progress_callback
)
elif isinstance(media, MessageMediaDocument):
elif isinstance(media, (MessageMediaDocument, Document)):
return self._download_document(
media, file, date, progress_callback
)
@ -1319,7 +1319,6 @@ class TelegramClient(TelegramBareClient):
def _download_photo(self, photo, file, date, progress_callback):
"""Specialized version of .download_media() for photos"""
# Determine the photo and its largest size
if isinstance(photo, MessageMediaPhoto):
photo = photo.photo
@ -1345,9 +1344,13 @@ class TelegramClient(TelegramBareClient):
)
return file
def _download_document(self, mm_doc, file, date, progress_callback):
def _download_document(self, document, file, date, progress_callback):
"""Specialized version of .download_media() for documents"""
document = mm_doc.document
if isinstance(document, MessageMediaDocument):
document = document.document
if not isinstance(document, Document):
return
file_size = document.size
possible_names = []
@ -1361,7 +1364,7 @@ class TelegramClient(TelegramBareClient):
))
file = self._get_proper_filename(
file, 'document', utils.get_extension(mm_doc),
file, 'document', utils.get_extension(document),
date=date, possible_names=possible_names
)

View File

@ -61,13 +61,13 @@ def get_extension(media):
# Documents will come with a mime type
if isinstance(media, MessageMediaDocument):
if isinstance(media.document, Document):
if media.document.mime_type == 'application/octet-stream':
# Octet stream are just bytes, which have no default extension
return ''
else:
extension = guess_extension(media.document.mime_type)
return extension if extension else ''
media = media.document
if isinstance(media, Document):
if media.mime_type == 'application/octet-stream':
# Octet stream are just bytes, which have no default extension
return ''
else:
return guess_extension(media.mime_type) or ''
return ''