Make .download_profile_photo() more friendly

This commit is contained in:
Lonami Exo 2017-08-23 01:35:12 +02:00
parent 9d15008f09
commit 8ed77418bf
2 changed files with 68 additions and 21 deletions

View File

@ -40,11 +40,11 @@ from .tl.functions import PingRequest
# All the types we need to work with # All the types we need to work with
from .tl.types import ( from .tl.types import (
ChatPhotoEmpty, DocumentAttributeAudio, DocumentAttributeFilename, DocumentAttributeAudio, DocumentAttributeFilename,
InputDocumentFileLocation, InputFileLocation, InputDocumentFileLocation, InputFileLocation,
InputMediaUploadedDocument, InputMediaUploadedPhoto, InputPeerEmpty, InputMediaUploadedDocument, InputMediaUploadedPhoto, InputPeerEmpty,
Message, MessageMediaContact, MessageMediaDocument, MessageMediaPhoto, Message, MessageMediaContact, MessageMediaDocument, MessageMediaPhoto,
UserProfilePhotoEmpty, InputUserSelf) InputUserSelf, UserProfilePhoto, ChatPhoto)
from .utils import find_user_or_chat, get_extension from .utils import find_user_or_chat, get_extension
@ -564,26 +564,72 @@ class TelegramClient(TelegramBareClient):
# region Downloading media requests # region Downloading media requests
def download_profile_photo(self, def download_profile_photo(self,
profile_photo, entity,
file_path, file_path='',
add_extension=True, add_extension=True,
download_big=True): download_big=True):
"""Downloads the profile photo for an user or a chat (including channels). """Downloads the profile photo for an user or a chat (channels too).
Returns False if no photo was provided, or if it was Empty""" Returns None if no photo was provided, or if it was Empty.
# TODO Support specifying an entity
if (not profile_photo or If an entity itself (an user, chat or channel) is given, the photo
isinstance(profile_photo, UserProfilePhotoEmpty) or to be downloaded will be downloaded automatically.
isinstance(profile_photo, ChatPhotoEmpty)):
return False On success, the file path is returned since it may differ.
The entity may be a phone or an username at the expense of
some performance loss.
"""
if not isinstance(entity, TLObject) or type(entity).subclass_of_id in (
0x2da17977, 0xc5af5d94, 0x1f4661b9, 0xd49a2697
):
# Maybe it is an user or a chat? Or their full versions?
#
# The hexadecimal numbers above are simply:
# hex(crc32(x.encode('ascii'))) for x in
# ('User', 'Chat', 'UserFull', 'ChatFull')
entity = self._get_entity(entity)
if not hasattr(entity, 'photo'):
# Special case: may be a ChatFull with photo:Photo
# This is different from a normal UserProfilePhoto and Chat
if hasattr(entity, 'chat_photo'):
return self._download_photo(
entity.chat_photo, file_path,
add_extension=add_extension
)
else:
# Give up
return None
was_dir = None
if os.path.isdir(file_path):
was_dir = file_path
file_path = '' # File path needs to be empty to infer it
if not file_path:
for attr in ('username', 'first_name', 'title'):
file_path = getattr(entity, attr, '')
if file_path:
break
if was_dir:
file_path = os.path.join(was_dir, file_path)
entity = entity.photo
if not isinstance(entity, UserProfilePhoto) and \
not isinstance(entity, ChatPhoto):
return None
if os.path.isdir(file_path) or not file_path:
file_path = 'profile_photo_{}'.format(datetime.now())
if add_extension: if add_extension:
file_path += get_extension(profile_photo) file_path += get_extension(entity)
if download_big: if download_big:
photo_location = profile_photo.photo_big photo_location = entity.photo_big
else: else:
photo_location = profile_photo.photo_small photo_location = entity.photo_small
# 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(
@ -594,7 +640,7 @@ class TelegramClient(TelegramBareClient):
), ),
file_path file_path
) )
return True return file_path
def download_media(self, def download_media(self,
message, message,

View File

@ -206,12 +206,13 @@ class InteractiveTelegramClient(TelegramClient):
# Download profile photo # Download profile photo
elif msg == '!dp': elif msg == '!dp':
output = str('usermedia/propic_{}'.format(entity.id)) print('Downloading profile picture to usermedia/...')
print('Downloading profile picture...') os.makedirs('usermedia', exist_ok=True)
success = self.download_profile_photo(entity.photo, output) output = self.download_profile_photo(entity, 'usermedia')
if success: if output:
print('Profile picture downloaded to {}'.format( print(
output)) 'Profile picture downloaded to {}'.format(output)
)
else: else:
print('No profile picture found for this user.') print('No profile picture found for this user.')