All download_* methods now accept streams

This commit is contained in:
hnikaein 2017-07-20 12:07:19 +04:30 committed by Lonami
parent 3be995b5a3
commit c9e566342e

View File

@ -542,33 +542,33 @@ class TelegramClient(TelegramBareClient):
def download_msg_media(self, def download_msg_media(self,
message_media, message_media,
file_path, file,
add_extension=True, add_extension=True,
progress_callback=None): progress_callback=None):
"""Downloads the given MessageMedia (Photo, Document or Contact) """Downloads the given MessageMedia (Photo, Document or Contact)
into the desired file_path, optionally finding its extension automatically into the desired file(a stream or str), optionally finding its extension automatically
The progress_callback should be a callback function which takes two parameters, The progress_callback should be a callback function which takes two parameters,
uploaded size (in bytes) and total file size (in bytes). uploaded size (in bytes) and total file size (in bytes).
This will be called every time a part is downloaded""" This will be called every time a part is downloaded"""
if type(message_media) == MessageMediaPhoto: if type(message_media) == MessageMediaPhoto:
return self.download_photo(message_media, file_path, add_extension, return self.download_photo(message_media, file, add_extension,
progress_callback) progress_callback)
elif type(message_media) == MessageMediaDocument: elif type(message_media) == MessageMediaDocument:
return self.download_document(message_media, file_path, return self.download_document(message_media, file,
add_extension, progress_callback) add_extension, progress_callback)
elif type(message_media) == MessageMediaContact: elif type(message_media) == MessageMediaContact:
return self.download_contact(message_media, file_path, return self.download_contact(message_media, file,
add_extension) add_extension)
def download_photo(self, def download_photo(self,
message_media_photo, message_media_photo,
file_path, file,
add_extension=False, add_extension=False,
progress_callback=None): progress_callback=None):
"""Downloads MessageMediaPhoto's largest size into the desired """Downloads MessageMediaPhoto's largest size into the desired
file_path, optionally finding its extension automatically file(a stream or str), optionally finding its extension automatically
The progress_callback should be a callback function which takes two parameters, The progress_callback should be a callback function which takes two parameters,
uploaded size (in bytes) and total file size (in bytes). uploaded size (in bytes) and total file size (in bytes).
This will be called every time a part is downloaded""" This will be called every time a part is downloaded"""
@ -579,8 +579,8 @@ class TelegramClient(TelegramBareClient):
file_size = largest_size.size file_size = largest_size.size
largest_size = largest_size.location largest_size = largest_size.location
if add_extension: if isinstance(file, str) and add_extension:
file_path += get_extension(message_media_photo) file += get_extension(message_media_photo)
# Download the media with the largest size input file location # Download the media with the largest size input file location
self.download_file( self.download_file(
@ -589,19 +589,19 @@ class TelegramClient(TelegramBareClient):
local_id=largest_size.local_id, local_id=largest_size.local_id,
secret=largest_size.secret secret=largest_size.secret
), ),
file_path, file,
file_size=file_size, file_size=file_size,
progress_callback=progress_callback progress_callback=progress_callback
) )
return file_path return file
def download_document(self, def download_document(self,
message_media_document, message_media_document,
file_path=None, file=None,
add_extension=True, add_extension=True,
progress_callback=None): progress_callback=None):
"""Downloads the given MessageMediaDocument into the desired """Downloads the given MessageMediaDocument into the desired
file_path, optionally finding its extension automatically. file(a stream or str), optionally finding its extension automatically.
If no file_path is given, it will try to be guessed from the document If no file_path is given, it will try to be guessed from the document
The progress_callback should be a callback function which takes two parameters, The progress_callback should be a callback function which takes two parameters,
uploaded size (in bytes) and total file size (in bytes). uploaded size (in bytes) and total file size (in bytes).
@ -610,21 +610,21 @@ class TelegramClient(TelegramBareClient):
file_size = document.size file_size = document.size
# If no file path was given, try to guess it from the attributes # If no file path was given, try to guess it from the attributes
if file_path is None: if file is None:
for attr in document.attributes: for attr in document.attributes:
if type(attr) == DocumentAttributeFilename: if type(attr) == DocumentAttributeFilename:
file_path = attr.file_name file = attr.file_name
break # This attribute has higher preference break # This attribute has higher preference
elif type(attr) == DocumentAttributeAudio: elif type(attr) == DocumentAttributeAudio:
file_path = '{} - {}'.format(attr.performer, attr.title) file = '{} - {}'.format(attr.performer, attr.title)
if file_path is None: if file is None:
raise ValueError('Could not infer a file_path for the document' raise ValueError('Could not infer a file_path for the document'
'. Please provide a valid file_path manually') '. Please provide a valid file_path manually')
if add_extension: if isinstance(file, str) and add_extension:
file_path += get_extension(message_media_document) file += get_extension(message_media_document)
self.download_file( self.download_file(
InputDocumentFileLocation( InputDocumentFileLocation(
@ -632,29 +632,40 @@ class TelegramClient(TelegramBareClient):
access_hash=document.access_hash, access_hash=document.access_hash,
version=document.version version=document.version
), ),
file_path, file,
file_size=file_size, file_size=file_size,
progress_callback=progress_callback progress_callback=progress_callback
) )
return file_path return file
@staticmethod @staticmethod
def download_contact(message_media_contact, file_path, add_extension=True): def download_contact(message_media_contact, file, add_extension=True):
"""Downloads a media contact using the vCard 4.0 format""" """Downloads a media contact using the vCard 4.0 format"""
first_name = message_media_contact.first_name first_name = message_media_contact.first_name
last_name = message_media_contact.last_name last_name = message_media_contact.last_name
phone_number = message_media_contact.phone_number phone_number = message_media_contact.phone_number
if isinstance(file, str):
# The only way we can save a contact in an understandable # The only way we can save a contact in an understandable
# way by phones is by using the .vCard format # way by phones is by using the .vCard format
if add_extension: if add_extension:
file_path += '.vcard' file += '.vcard'
# Ensure that we'll be able to download the contact # Ensure that we'll be able to download the contact
utils.ensure_parent_dir_exists(file_path) utils.ensure_parent_dir_exists(file)
with open(file_path, 'w', encoding='utf-8') as file: with open(file, 'w', encoding='utf-8') as f:
f.write('BEGIN:VCARD\n')
f.write('VERSION:4.0\n')
f.write('N:{};{};;;\n'.format(first_name, last_name
if last_name else ''))
f.write('FN:{}\n'.format(' '.join((first_name, last_name))))
f.write('TEL;TYPE=cell;VALUE=uri:tel:+{}\n'.format(
phone_number))
f.write('END:VCARD\n')
else:
file.write('BEGIN:VCARD\n') file.write('BEGIN:VCARD\n')
file.write('VERSION:4.0\n') file.write('VERSION:4.0\n')
file.write('N:{};{};;;\n'.format(first_name, last_name file.write('N:{};{};;;\n'.format(first_name, last_name
@ -664,7 +675,7 @@ class TelegramClient(TelegramBareClient):
phone_number)) phone_number))
file.write('END:VCARD\n') file.write('END:VCARD\n')
return file_path return file
# endregion # endregion