From b98ec3efad66d4357057c22e71e416cbb5a8f9b9 Mon Sep 17 00:00:00 2001 From: darksidecat <58224121+darksidecat@users.noreply.github.com> Date: Sun, 7 Nov 2021 01:28:12 +0200 Subject: [PATCH] add `Bot.download_file` aliases ability to save files to a directory and automatically create directories (#694) * add destination_dir and make_dirs parameters to bot download aliases * add the ability to save files to a directory with path completion based on file_path, * add an option to automatically create directories in the file path * Downloadable mixin uses directory creation parameter in bot methods --- aiogram/bot/base.py | 36 ++++++++--- aiogram/bot/bot.py | 23 +++++-- aiogram/types/mixins.py | 7 +-- tests/test_bot/test_bot_download_file.py | 78 ++++++++++++++++++++++++ 4 files changed, 127 insertions(+), 17 deletions(-) create mode 100644 tests/test_bot/test_bot_download_file.py diff --git a/aiogram/bot/base.py b/aiogram/bot/base.py index 07e44c1c..1e9bcc15 100644 --- a/aiogram/bot/base.py +++ b/aiogram/bot/base.py @@ -1,6 +1,8 @@ import asyncio import contextlib import io +import os +import pathlib import ssl import typing import warnings @@ -208,28 +210,48 @@ class BaseBot: return await api.make_request(self.session, self.server, self.__token, method, data, files, proxy=self.proxy, proxy_auth=self.proxy_auth, timeout=self.timeout, **kwargs) - async def download_file(self, file_path: base.String, - destination: Optional[base.InputFile] = None, - timeout: Optional[base.Integer] = sentinel, - chunk_size: Optional[base.Integer] = 65536, - seek: Optional[base.Boolean] = True) -> Union[io.BytesIO, io.FileIO]: + async def download_file( + self, + file_path: base.String, + destination: Optional[Union[base.InputFile, pathlib.Path]] = None, + timeout: Optional[base.Integer] = sentinel, + chunk_size: Optional[base.Integer] = 65536, + seek: Optional[base.Boolean] = True, + destination_dir: Optional[Union[str, pathlib.Path]] = None, + make_dirs: Optional[base.Boolean] = True, + ) -> Union[io.BytesIO, io.FileIO]: """ - Download file by file_path to destination + Download file by file_path to destination file or directory if You want to automatically create destination (:class:`io.BytesIO`) use default value of destination and handle result of this method. + At most one of these parameters can be used: :param destination:, :param destination_dir: + :param file_path: file path on telegram server (You can get it from :obj:`aiogram.types.File`) :type file_path: :obj:`str` :param destination: filename or instance of :class:`io.IOBase`. For e. g. :class:`io.BytesIO` :param timeout: Integer :param chunk_size: Integer :param seek: Boolean - go to start of file when downloading is finished. + :param destination_dir: directory for saving files + :param make_dirs: Make dirs if not exist :return: destination """ - if destination is None: + if destination and destination_dir: + raise ValueError( + "Use only one of the parameters:destination or destination_dir." + ) + + if destination is None and destination_dir is None: destination = io.BytesIO() + elif destination_dir: + destination = os.path.join(destination_dir, file_path) + + if make_dirs and not isinstance(destination, io.IOBase) and os.path.dirname(destination): + os.makedirs(os.path.dirname(destination), exist_ok=True) + url = self.get_file_url(file_path) dest = destination if isinstance(destination, io.IOBase) else open(destination, 'wb') diff --git a/aiogram/bot/bot.py b/aiogram/bot/bot.py index 304b000d..22b1c91c 100644 --- a/aiogram/bot/bot.py +++ b/aiogram/bot/bot.py @@ -1,6 +1,7 @@ from __future__ import annotations import datetime +import pathlib import typing import warnings @@ -43,25 +44,37 @@ class Bot(BaseBot, DataMixin, ContextInstanceMixin): if hasattr(self, '_me'): delattr(self, '_me') - async def download_file_by_id(self, file_id: base.String, destination=None, - timeout: base.Integer = 30, chunk_size: base.Integer = 65536, - seek: base.Boolean = True): + async def download_file_by_id( + self, + file_id: base.String, + destination: typing.Optional[base.InputFile, pathlib.Path] = None, + timeout: base.Integer = 30, + chunk_size: base.Integer = 65536, + seek: base.Boolean = True, + destination_dir: typing.Optional[typing.Union[str, pathlib.Path]] = None, + make_dirs: typing.Optional[base.Boolean] = True, + ): """ - Download file by file_id to destination + Download file by file_id to destination file or directory if You want to automatically create destination (:class:`io.BytesIO`) use default value of destination and handle result of this method. + At most one of these parameters can be used: :param destination:, :param destination_dir: + :param file_id: str :param destination: filename or instance of :class:`io.IOBase`. For e. g. :class:`io.BytesIO` :param timeout: int :param chunk_size: int :param seek: bool - go to start of file when downloading is finished + :param destination_dir: directory for saving files + :param make_dirs: Make dirs if not exist :return: destination """ file = await self.get_file(file_id) return await self.download_file(file_path=file.file_path, destination=destination, - timeout=timeout, chunk_size=chunk_size, seek=seek) + timeout=timeout, chunk_size=chunk_size, seek=seek, + destination_dir=destination_dir, make_dirs=make_dirs) # === Getting updates === # https://core.telegram.org/bots/api#getting-updates diff --git a/aiogram/types/mixins.py b/aiogram/types/mixins.py index 83c65032..7d06d4c4 100644 --- a/aiogram/types/mixins.py +++ b/aiogram/types/mixins.py @@ -49,7 +49,6 @@ class Downloadable: destination, destination_dir, destination_file, - make_dirs ) return await self.bot.download_file( @@ -58,9 +57,10 @@ class Downloadable: timeout=timeout, chunk_size=chunk_size, seek=seek, + make_dirs=make_dirs ) - async def _prepare_destination(self, dest, destination_dir, destination_file, make_dirs): + async def _prepare_destination(self, dest, destination_dir, destination_file): file = await self.get_file() if not(any((dest, destination_dir, destination_file))): @@ -87,9 +87,6 @@ class Downloadable: else: raise TypeError("destination_file must be str, pathlib.Path or io.IOBase type") - if make_dirs and os.path.dirname(destination): - os.makedirs(os.path.dirname(destination), exist_ok=True) - return file, destination async def get_file(self): diff --git a/tests/test_bot/test_bot_download_file.py b/tests/test_bot/test_bot_download_file.py new file mode 100644 index 00000000..75710fcc --- /dev/null +++ b/tests/test_bot/test_bot_download_file.py @@ -0,0 +1,78 @@ +import os +from io import BytesIO +from pathlib import Path + +import pytest + +from aiogram import Bot +from aiogram.types import File +from tests import TOKEN +from tests.types.dataset import FILE + +pytestmark = pytest.mark.asyncio + + +@pytest.fixture(name='bot') +async def bot_fixture(): + async def get_file(): + return File(**FILE) + + """ Bot fixture """ + _bot = Bot(TOKEN) + _bot.get_file = get_file + yield _bot + await _bot.session.close() + + +@pytest.fixture +def file(): + return File(**FILE) + + +@pytest.fixture +def tmppath(tmpdir, request): + os.chdir(tmpdir) + yield Path(tmpdir) + os.chdir(request.config.invocation_dir) + + +class TestBotDownload: + async def test_download_file(self, tmppath, bot, file): + f = await bot.download_file(file_path=file.file_path) + assert len(f.read()) != 0 + + async def test_download_file_destination(self, tmppath, bot, file): + await bot.download_file(file_path=file.file_path, destination="test.file") + assert os.path.isfile(tmppath.joinpath('test.file')) + + async def test_download_file_destination_with_dir(self, tmppath, bot, file): + await bot.download_file(file_path=file.file_path, + destination=os.path.join('dir_name', 'file_name')) + assert os.path.isfile(tmppath.joinpath('dir_name', 'file_name')) + + async def test_download_file_destination_raise_file_not_found(self, tmppath, bot, file): + with pytest.raises(FileNotFoundError): + await bot.download_file(file_path=file.file_path, + destination=os.path.join('dir_name', 'file_name'), + make_dirs=False) + + async def test_download_file_destination_io_bytes(self, tmppath, bot, file): + f = BytesIO() + await bot.download_file(file_path=file.file_path, + destination=f) + assert len(f.read()) != 0 + + async def test_download_file_raise_value_error(self, tmppath, bot, file): + with pytest.raises(ValueError): + await bot.download_file(file_path=file.file_path, destination="a", destination_dir="b") + + async def test_download_file_destination_dir(self, tmppath, bot, file): + await bot.download_file(file_path=file.file_path, destination_dir='test_dir') + assert os.path.isfile(tmppath.joinpath('test_dir', file.file_path)) + + async def test_download_file_destination_dir_raise_file_not_found(self, tmppath, bot, file): + with pytest.raises(FileNotFoundError): + await bot.download_file(file_path=file.file_path, + destination_dir='test_dir', + make_dirs=False) + assert os.path.isfile(tmppath.joinpath('test_dir', file.file_path))