From 753396330d9f2995def7f95dcf30f544a4dd9e85 Mon Sep 17 00:00:00 2001 From: Alex Root Junior Date: Tue, 21 Nov 2017 20:31:35 +0200 Subject: [PATCH] Implemented new one InputFile interface for sending local files. --- aiogram/bot/api.py | 3 + aiogram/types/base.py | 2 +- aiogram/types/input_file.py | 146 +++++++++++++++++++++++++++++++++-- aiogram/types/input_media.py | 29 ++----- 4 files changed, 149 insertions(+), 31 deletions(-) diff --git a/aiogram/bot/api.py b/aiogram/bot/api.py index 6c7b19c5..9025ee00 100644 --- a/aiogram/bot/api.py +++ b/aiogram/bot/api.py @@ -5,6 +5,7 @@ from http import HTTPStatus import aiohttp +from .. import types from ..utils import exceptions from ..utils import json from ..utils.helper import Helper, HelperMode, Item @@ -113,6 +114,8 @@ def _compose_data(params=None, files=None): filename, fileobj = f else: raise ValueError('Tuple must have exactly 2 elements: filename, fileobj') + elif isinstance(f, types.InputFile): + filename, fileobj = f.get_filename(), f.get_file() else: filename, fileobj = _guess_filename(f) or key, f diff --git a/aiogram/types/base.py b/aiogram/types/base.py index 03afcb53..da802ae3 100644 --- a/aiogram/types/base.py +++ b/aiogram/types/base.py @@ -13,7 +13,7 @@ ALIASES_ATTR_NAME = '_aliases' __all__ = ('MetaTelegramObject', 'TelegramObject') # Binding of builtin types -InputFile = TypeVar('InputFile', io.BytesIO, io.FileIO, str) +InputFile = TypeVar('InputFile', 'InputFile', io.BytesIO, io.FileIO, str) String = TypeVar('String', bound=str) Integer = TypeVar('Integer', bound=int) Float = TypeVar('Float', bound=float) diff --git a/aiogram/types/input_file.py b/aiogram/types/input_file.py index 716275c3..3c69b26a 100644 --- a/aiogram/types/input_file.py +++ b/aiogram/types/input_file.py @@ -1,7 +1,15 @@ +import io +import logging +import os +import tempfile +import time + +import aiohttp + from . import base +from ..bot import api - -# TODO: Interface for sending files +log = logging.getLogger('aiogram') class InputFile(base.TelegramObject): @@ -9,12 +17,134 @@ class InputFile(base.TelegramObject): This object represents the contents of a file to be uploaded. Must be posted using multipart/form-data in the usual way that files are uploaded via the browser. + Also that is not typical TelegramObject! + https://core.telegram.org/bots/api#inputfile """ - def __init__(self, file_id=None, path=None, url=None, filename=None): - self.file_id = file_id - self.path = path - self.url = url - self.filename = filename - super(InputFile, self).__init__() + def __init__(self, path_or_bytesio, filename=None, conf=None): + """ + + :param path_or_bytesio: + :param filename: + :param conf: + """ + super(InputFile, self).__init__(conf=conf) + if isinstance(path_or_bytesio, str): + # As path + self._file = open(path_or_bytesio, 'rb') + self._path = path_or_bytesio + if filename is None: + filename = os.path.split(path_or_bytesio)[-1] + else: + # As io.BytesIO + assert isinstance(path_or_bytesio, io.IOBase) + self._path = None + self._file = path_or_bytesio + + self._filename = filename + + def __del__(self): + """ + Close file descriptor + """ + if not hasattr(self, '_file'): + return + self._file.close() + del self._file + + if self.conf.get('downloaded') and self.conf.get('temp'): + log.debug(f"Unlink file '{self._path}'") + os.unlink(self._path) + + def get_filename(self) -> str: + """ + Get file name + + :return: name + """ + if self._filename is None: + self._filename = api._guess_filename(self._file) + return self._filename + + def get_file(self): + """ + Get file object + + :return: + """ + return self._file + + @classmethod + async def from_url(cls, url, filename=None, temp_file=False, chunk_size=65536): + """ + Download file from URL + + Manually is not required action. You can send urls instead! + + :param url: target URL + :param filename: optional. set custom file name + :param temp_file: use temporary file + :param chunk_size: + + :return: InputFile + """ + conf = { + 'downloaded': True, + 'url': url + } + + # Let's do magic with the filename + if filename: + filename_prefix, _, ext = filename.rpartition('.') + file_suffix = '.' + ext if ext else '' + else: + filename_prefix, _, ext = url.rpartition('/')[-1].rpartition('.') + file_suffix = '.' + ext if ext else '' + filename = filename_prefix + file_suffix + + async with aiohttp.ClientSession() as session: + start = time.time() + async with session.get(url) as response: + if temp_file: + # Create temp file + fd, path = tempfile.mkstemp(suffix=file_suffix, prefix=filename_prefix + '_') + file = conf['temp'] = path + + # Save file in temp directory + with open(fd, 'wb') as f: + await cls._process_stream(response, f, chunk_size=chunk_size) + else: + # Save file in memory + file = await cls._process_stream(response, io.BytesIO(), chunk_size=chunk_size) + + log.debug(f"File successful downloaded at {round(time.time() - start, 2)} seconds from '{url}'") + return cls(file, filename, conf=conf) + + @classmethod + async def _process_stream(cls, response, writer, chunk_size=65536): + """ + Transfer data + + :param response: + :param writer: + :param chunk_size: + :return: + """ + while True: + chunk = await response.content.read(chunk_size) + if not chunk: + break + writer.write(chunk) + + if writer.seekable(): + writer.seek(0) + + return writer + + def to_python(self): + raise TypeError('Object of this type is not exportable!') + + @classmethod + def to_object(cls, data): + raise TypeError('Object of this type is not importable!') diff --git a/aiogram/types/input_media.py b/aiogram/types/input_media.py index 98a7bc86..391e6581 100644 --- a/aiogram/types/input_media.py +++ b/aiogram/types/input_media.py @@ -4,6 +4,7 @@ import typing from . import base from . import fields +from .input_file import InputFile ATTACHMENT_PREFIX = 'attach://' @@ -28,20 +29,6 @@ class InputMedia(base.TelegramObject): @file.setter def file(self, file: io.IOBase): - if self.conf.get('cache'): - # File must be not closed before sending media. - # Read file into BytesIO - if isinstance(file, io.BufferedIOBase): - # Go to start of file - if file.seekable(): - file.seek(0) - # Read - temp_file = io.BytesIO(file.read()) - # Reset cursor - file.seek(0) - # Replace variable - file = temp_file - setattr(self, '_file', file) self.media = ATTACHMENT_PREFIX + secrets.token_urlsafe(16) @@ -59,10 +46,10 @@ class InputMediaPhoto(InputMedia): https://core.telegram.org/bots/api#inputmediaphoto """ - def __init__(self, media: base.InputFile, caption: base.String = None, cache=True): - super(InputMediaPhoto, self).__init__(type='photo', media=media, caption=caption, conf={'cache': cache}) + def __init__(self, media: base.InputFile, caption: base.String = None): + super(InputMediaPhoto, self).__init__(type='photo', media=media, caption=caption) - if isinstance(media, io.IOBase): + if isinstance(media, (io.IOBase, InputFile)): self.file = media @@ -77,13 +64,11 @@ class InputMediaVideo(InputMedia): duration: base.Integer = fields.Field() def __init__(self, media: base.InputFile, caption: base.String = None, - width: base.Integer = None, height: base.Integer = None, duration: base.Integer = None, - cache=True): + width: base.Integer = None, height: base.Integer = None, duration: base.Integer = None): super(InputMediaVideo, self).__init__(type='video', media=media, caption=caption, - width=width, height=height, duration=duration, - conf={'cache': cache}) + width=width, height=height, duration=duration) - if isinstance(media, io.IOBase): + if isinstance(media, (io.IOBase, InputFile)): self.file = media