Implemented new one InputFile interface for sending local files.

This commit is contained in:
Alex Root Junior 2017-11-21 20:31:35 +02:00
parent 751929cb1d
commit 753396330d
4 changed files with 149 additions and 31 deletions

View file

@ -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

View file

@ -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)

View file

@ -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!')

View file

@ -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