From a42effc2b5c274f7ef8d7cca5fbe854eacf63598 Mon Sep 17 00:00:00 2001 From: Lonami Date: Sun, 9 Oct 2016 12:57:38 +0200 Subject: [PATCH] Externalized TelegramClient utilities and added more --- telethon/interactive_telegram_client.py | 11 ++- telethon/telegram_client.py | 100 ++++-------------------- telethon/utils/__init__.py | 1 + telethon/utils/tl_utils.py | 89 +++++++++++++++++++++ 4 files changed, 114 insertions(+), 87 deletions(-) create mode 100644 telethon/utils/tl_utils.py diff --git a/telethon/interactive_telegram_client.py b/telethon/interactive_telegram_client.py index b01c613b..5fa8dabf 100644 --- a/telethon/interactive_telegram_client.py +++ b/telethon/interactive_telegram_client.py @@ -1,6 +1,9 @@ from telethon.tl.types import UpdateShortChatMessage from telethon.tl.types import UpdateShortMessage from telethon import TelegramClient + +from telethon.utils import get_display_name, get_input_peer + import shutil # Get the (current) number of lines in the terminal @@ -71,7 +74,7 @@ class InteractiveTelegramClient(TelegramClient): # Display them so the user can choose for i, entity in enumerate(entities): i += 1 # 1-based index for normies - print('{}. {}'.format(i, self.get_display_name(entity))) + print('{}. {}'.format(i, get_display_name(entity))) # Let the user decide who they want to talk to print() @@ -97,10 +100,10 @@ class InteractiveTelegramClient(TelegramClient): # Retrieve the selected user (or chat, or channel) entity = entities[i] - input_peer = self.get_input_peer(entity) + input_peer = get_input_peer(entity) # Show some information - print_title('Chat with "{}"'.format(self.get_display_name(entity))) + print_title('Chat with "{}"'.format(get_display_name(entity))) print('Available commands:') print(' !q: Quits the current chat.') print(' !Q: Quits the current chat and exits.') @@ -166,7 +169,7 @@ class InteractiveTelegramClient(TelegramClient): print('Profile picture downloaded to {}'.format(output)) else: print('"{}" does not seem to have a profile picture.' - .format(self.get_display_name(entity))) + .format(get_display_name(entity))) # Send chat message (if any) elif msg: diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py index 28be68fd..bdd90a38 100644 --- a/telethon/telegram_client.py +++ b/telethon/telegram_client.py @@ -2,7 +2,7 @@ import platform from datetime import datetime, timedelta from hashlib import md5 from os import path, listdir -from mimetypes import guess_extension, guess_type +from mimetypes import guess_type # For sending and receiving requests from telethon.tl import MTProtoRequest @@ -10,15 +10,6 @@ from telethon.tl import Session # The Requests and types that we'll be using from telethon.tl.functions.upload import SaveBigFilePartRequest -from telethon.tl.types import \ - User, Chat, Channel, \ - PeerUser, PeerChat, PeerChannel, \ - InputPeerUser, InputPeerChat, InputPeerChannel, InputPeerEmpty, \ - UserProfilePhotoEmpty, ChatPhotoEmpty, \ - InputFile, InputFileLocation, InputMediaUploadedPhoto, InputMediaUploadedDocument, \ - MessageMediaContact, MessageMediaDocument, MessageMediaPhoto, \ - DocumentAttributeAudio, DocumentAttributeFilename, InputDocumentFileLocation - from telethon.tl.functions import InvokeWithLayerRequest, InitConnectionRequest from telethon.tl.functions.help import GetConfigRequest from telethon.tl.functions.auth import SendCodeRequest, SignInRequest, SignUpRequest, LogOutRequest @@ -28,8 +19,18 @@ from telethon.tl.functions.messages import \ SendMessageRequest, SendMediaRequest, \ ReadHistoryRequest +# All the types we need to work with +from telethon.tl.types import \ + InputPeerEmpty, \ + UserProfilePhotoEmpty, ChatPhotoEmpty, \ + InputFile, InputFileLocation, InputMediaUploadedPhoto, InputMediaUploadedDocument, \ + MessageMediaContact, MessageMediaDocument, MessageMediaPhoto, \ + DocumentAttributeAudio, DocumentAttributeFilename, InputDocumentFileLocation + +# Import some externalized utilities to work with the Telegram types and more import telethon.helpers as utils import telethon.network.authenticator as authenticator +from telethon.utils import find_user_or_chat, get_appropiate_part_size, get_extension from telethon.errors import * from telethon.network import MtProtoSender, TcpTransport @@ -227,7 +228,7 @@ class TelegramClient: offset_peer=offset_peer, limit=count)) return (r.dialogs, - [self.find_user_or_chat(d.peer, r.users, r.chats) for d in r.dialogs]) + [find_user_or_chat(d.peer, r.users, r.chats) for d in r.dialogs]) # endregion @@ -325,7 +326,7 @@ class TelegramClient: """ file_size = path.getsize(file_path) if not part_size_kb: - part_size_kb = self.find_appropiate_part_size(file_size) + part_size_kb = get_appropiate_part_size(file_size) if part_size_kb > 512: raise ValueError('The part size must be less or equal to 512KB') @@ -422,9 +423,8 @@ class TelegramClient: isinstance(profile_photo, ChatPhotoEmpty)): return False - # Photos are always compressed into a .jpg by Telegram if add_extension: - file_path += '.jpg' + file_path += get_extension(profile_photo) if download_big: photo_location = profile_photo.photo_big @@ -467,9 +467,8 @@ class TelegramClient: file_size = largest_size.size largest_size = largest_size.location - # Photos are always compressed into a .jpg by Telegram if add_extension: - file_path += '.jpg' + file_path += get_extension(message_media_photo) # Download the media with the largest size input file location self.download_file_loc(InputFileLocation(volume_id=largest_size.volume_id, @@ -502,11 +501,8 @@ class TelegramClient: if file_path is None: print('Could not determine a filename for the document') - # Guess the extension based on the mime_type if add_extension: - ext = guess_extension(document.mime_type) - if ext is not None: - file_path += ext + file_path += get_extension(document.mime_type) self.download_file_loc(InputDocumentFileLocation(id=document.id, access_hash=document.access_hash, @@ -551,7 +547,7 @@ class TelegramClient: if not file_size: raise ValueError('A part size value must be provided') else: - part_size_kb = self.find_appropiate_part_size(file_size) + part_size_kb = get_appropiate_part_size(file_size) part_size = int(part_size_kb * 1024) if part_size % 1024 != 0: @@ -582,68 +578,6 @@ class TelegramClient: # endregion - # region Utilities - - @staticmethod - def get_display_name(entity): - """Gets the input peer for the given "entity" (user, chat or channel) - Returns None if it was not found""" - if isinstance(entity, User): - if entity.last_name is not None: - return '{} {}'.format(entity.first_name, entity.last_name) - return entity.first_name - - if isinstance(entity, Chat) or isinstance(entity, Channel): - return entity.title - - @staticmethod - def get_input_peer(entity): - """Gets the input peer for the given "entity" (user, chat or channel). - Returns None if it was not found""" - if isinstance(entity, User): - return InputPeerUser(entity.id, entity.access_hash) - if isinstance(entity, Chat): - return InputPeerChat(entity.id) - if isinstance(entity, Channel): - return InputPeerChannel(entity.id, entity.access_hash) - - @staticmethod - def find_user_or_chat(peer, users, chats): - """Finds the corresponding user or chat given a peer. - Returns None if it was not found""" - try: - if isinstance(peer, PeerUser): - user = next(u for u in users if u.id == peer.user_id) - return user - - elif isinstance(peer, PeerChat): - chat = next(c for c in chats if c.id == peer.chat_id) - return chat - - elif isinstance(peer, PeerChannel): - channel = next(c for c in chats if c.id == peer.channel_id) - return channel - - except StopIteration: - return None - - @staticmethod - def find_appropiate_part_size(file_size): - if file_size <= 1048576: # 1MB - return 32 - if file_size <= 10485760: # 10MB - return 64 - if file_size <= 393216000: # 375MB - return 128 - if file_size <= 786432000: # 750MB - return 256 - if file_size <= 1572864000: # 1500MB - return 512 - - raise ValueError('File size too large') - - # endregion - # region Updates handling def add_update_handler(self, handler): diff --git a/telethon/utils/__init__.py b/telethon/utils/__init__.py index 4923827f..c8f99ae5 100755 --- a/telethon/utils/__init__.py +++ b/telethon/utils/__init__.py @@ -1,2 +1,3 @@ from .binary_writer import BinaryWriter from .binary_reader import BinaryReader +from .tl_utils import * diff --git a/telethon/utils/tl_utils.py b/telethon/utils/tl_utils.py new file mode 100644 index 00000000..d9a8a20a --- /dev/null +++ b/telethon/utils/tl_utils.py @@ -0,0 +1,89 @@ +"""Utilities for working with TLObjects. + We have these because some TLObjects differ in very little things, + for example, some may have an `user_id` attribute and other a `chat_id` but, + after all, both are the same attribute, IDs.""" +from mimetypes import guess_extension + +from telethon.tl.types import \ + User, Chat, Channel, \ + PeerUser, PeerChat, PeerChannel, \ + InputPeerUser, InputPeerChat, InputPeerChannel, \ + UserProfilePhoto, ChatPhoto, \ + MessageMediaPhoto, MessageMediaDocument + + +def get_display_name(entity): + """Gets the input peer for the given "entity" (user, chat or channel) + Returns None if it was not found""" + if isinstance(entity, User): + if entity.last_name is not None: + return '{} {}'.format(entity.first_name, entity.last_name) + return entity.first_name + + if isinstance(entity, Chat) or isinstance(entity, Channel): + return entity.title + + +def get_extension(media): + """Gets the corresponding extension for any Telegram media""" + + # Photos are always compressed as .jpg by Telegram + if (isinstance(media, UserProfilePhoto) or + isinstance(media, ChatPhoto) or + isinstance(media, MessageMediaPhoto)): + return '.jpg' + + # Documents will come with a mime type, from which we can guess their mime type + if isinstance(media, MessageMediaDocument): + extension = guess_extension(media.document.mime_type) + return extension if extension else '' + + return None + + +def get_input_peer(entity): + """Gets the input peer for the given "entity" (user, chat or channel). + Returns None if it was not found""" + if isinstance(entity, User): + return InputPeerUser(entity.id, entity.access_hash) + if isinstance(entity, Chat): + return InputPeerChat(entity.id) + if isinstance(entity, Channel): + return InputPeerChannel(entity.id, entity.access_hash) + + +def find_user_or_chat(peer, users, chats): + """Finds the corresponding user or chat given a peer. + Returns None if it was not found""" + try: + if isinstance(peer, PeerUser): + user = next(u for u in users if u.id == peer.user_id) + return user + + elif isinstance(peer, PeerChat): + chat = next(c for c in chats if c.id == peer.chat_id) + return chat + + elif isinstance(peer, PeerChannel): + channel = next(c for c in chats if c.id == peer.channel_id) + return channel + + except StopIteration: + return None + + +def get_appropiate_part_size(file_size): + """Gets the appropiate part size when uploading or downloading files, + given an initial file size""" + if file_size <= 1048576: # 1MB + return 32 + if file_size <= 10485760: # 10MB + return 64 + if file_size <= 393216000: # 375MB + return 128 + if file_size <= 786432000: # 750MB + return 256 + if file_size <= 1572864000: # 1500MB + return 512 + + raise ValueError('File size too large')