From 78ee7873101403edaa0466475fff1ab2fda1d111 Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Fri, 17 Jan 2020 11:11:10 +0100 Subject: [PATCH 1/2] Fix utils._get_extension not working in pathlib objects This was found while testing #1371. --- telethon/utils.py | 3 +++ tests/telethon/test_utils.py | 37 ++++++++++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/telethon/utils.py b/telethon/utils.py index 8b12759f..0c63d143 100644 --- a/telethon/utils.py +++ b/telethon/utils.py @@ -12,6 +12,7 @@ import logging import math import mimetypes import os +import pathlib import re import struct from collections import namedtuple @@ -741,6 +742,8 @@ def _get_extension(file): """ if isinstance(file, str): return os.path.splitext(file)[-1] + elif isinstance(file, pathlib.Path): + return file.suffix elif isinstance(file, bytes): kind = imghdr.what(io.BytesIO(file)) return ('.' + kind) if kind else '' diff --git a/tests/telethon/test_utils.py b/tests/telethon/test_utils.py index d6f97e68..89391e36 100644 --- a/tests/telethon/test_utils.py +++ b/tests/telethon/test_utils.py @@ -1,3 +1,6 @@ +import io +import pathlib + from telethon import utils from telethon.tl.types import ( MessageMediaGame, Game, PhotoEmpty @@ -16,3 +19,37 @@ def test_game_input_media_memory_error(): )) input_media = utils.get_input_media(media) bytes(input_media) # <- shouldn't raise `MemoryError` + + +def test_private_get_extension(): + # Positive cases + png_header = bytes.fromhex('89 50 4e 47 0d 0a 1a 0a 00 00 00 0d 49 48 44 52') + png_buffer = io.BytesIO(png_header) + + class CustomFd: + def __init__(self, name): + self.name = name + + assert utils._get_extension('foo.bar.baz') == '.baz' + assert utils._get_extension(pathlib.Path('foo.bar.baz')) == '.baz' + assert utils._get_extension(png_header) == '.png' + assert utils._get_extension(png_buffer) == '.png' + assert utils._get_extension(png_buffer) == '.png' # make sure it did seek back + assert utils._get_extension(CustomFd('foo.bar.baz')) == '.baz' + + # Negative cases + null_header = bytes.fromhex('00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00') + null_buffer = io.BytesIO(null_header) + + empty_header = bytes() + empty_buffer = io.BytesIO(empty_header) + + assert utils._get_extension('foo') == '' + assert utils._get_extension(pathlib.Path('foo')) == '' + assert utils._get_extension(null_header) == '' + assert utils._get_extension(null_buffer) == '' + assert utils._get_extension(null_buffer) == '' # make sure it did seek back + assert utils._get_extension(empty_header) == '' + assert utils._get_extension(empty_buffer) == '' + assert utils._get_extension(empty_buffer) == '' # make sure it did seek back + assert utils._get_extension(CustomFd('foo')) == '' From 76cc076d6124972f9d0041f567f9769b05d5b9a5 Mon Sep 17 00:00:00 2001 From: Alexhol Date: Fri, 17 Jan 2020 13:12:20 +0300 Subject: [PATCH 2/2] Fix send_file not considering videos for albums (#1371) --- telethon/client/uploads.py | 36 ++++++++++++++++++++++++------------ 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/telethon/client/uploads.py b/telethon/client/uploads.py index 22e1de67..a6879914 100644 --- a/telethon/client/uploads.py +++ b/telethon/client/uploads.py @@ -269,7 +269,7 @@ class UploadMethods: # First check if the user passed an iterable, in which case # we may want to send as an album if all are photo files. if utils.is_list_like(file): - image_captions = [] + media_captions = [] document_captions = [] if utils.is_list_like(caption): captions = caption @@ -277,28 +277,29 @@ class UploadMethods: captions = [caption] # TODO Fix progress_callback - images = [] + media = [] if force_document: documents = file else: documents = [] for doc, cap in itertools.zip_longest(file, captions): - if utils.is_image(doc): - images.append(doc) - image_captions.append(cap) + if utils.is_image(doc) or utils.is_video(doc): + media.append(doc) + media_captions.append(cap) else: documents.append(doc) document_captions.append(cap) result = [] - while images: + while media: result += await self._send_album( - entity, images[:10], caption=image_captions[:10], + entity, media[:10], caption=media_captions[:10], progress_callback=progress_callback, reply_to=reply_to, - parse_mode=parse_mode, silent=silent, schedule=schedule + parse_mode=parse_mode, silent=silent, schedule=schedule, + supports_streaming=supports_streaming ) - images = images[10:] - image_captions = image_captions[10:] + media = media[10:] + media_captions = media_captions[10:] for doc, cap in zip(documents, captions): result.append(await self.send_file( @@ -349,7 +350,8 @@ class UploadMethods: async def _send_album(self: 'TelegramClient', entity, files, caption='', progress_callback=None, reply_to=None, - parse_mode=(), silent=None, schedule=None): + parse_mode=(), silent=None, schedule=None, + supports_streaming=None): """Specialized version of .send_file for albums""" # We don't care if the user wants to avoid cache, we will use it # anyway. Why? The cached version will be exactly the same thing @@ -377,7 +379,8 @@ class UploadMethods: # :tl:`InputMediaUploadedPhoto`. However using that will # make it `raise MediaInvalidError`, so we need to upload # it as media and then convert that to :tl:`InputMediaPhoto`. - fh, fm, _ = await self._file_to_media(file) + fh, fm, _ = await self._file_to_media( + file, supports_streaming=supports_streaming) if isinstance(fm, types.InputMediaUploadedPhoto): r = await self(functions.messages.UploadMediaRequest( entity, media=fm @@ -386,6 +389,15 @@ class UploadMethods: fh.md5, fh.size, utils.get_input_photo(r.photo)) fm = utils.get_input_media(r.photo) + elif isinstance(fm, types.InputMediaUploadedDocument): + r = await self(functions.messages.UploadMediaRequest( + entity, media=fm + )) + self.session.cache_file( + fh.md5, fh.size, utils.get_input_document(r.document)) + + fm = utils.get_input_media( + r.document, supports_streaming=supports_streaming) if captions: caption, msg_entities = captions.pop()