diff --git a/aiogram/api/client/bot.py b/aiogram/api/client/bot.py index b58e6d76..37dad2e7 100644 --- a/aiogram/api/client/bot.py +++ b/aiogram/api/client/bot.py @@ -4,7 +4,17 @@ import datetime import io import pathlib from contextlib import asynccontextmanager -from typing import Any, AsyncGenerator, AsyncIterator, BinaryIO, List, Optional, TypeVar, Union +from typing import ( + Any, + AsyncGenerator, + AsyncIterator, + BinaryIO, + List, + Optional, + TypeVar, + Union, + cast, +) import aiofiles from async_lru import alru_cache @@ -112,6 +122,7 @@ from ..types import ( UserProfilePhotos, WebhookInfo, ) +from ..types.downloadable import Downloadable from .session.aiohttp import AiohttpSession from .session.base import BaseSession @@ -199,18 +210,15 @@ class Bot(ContextInstanceMixin["Bot"]): ) -> Optional[BinaryIO]: """ Download file by file_path to destination. + If you want to automatically create destination (:class:`io.BytesIO`) use default value of destination and handle result of this method. + :param file_path: File path on Telegram server (You can get it from :obj:`aiogram.types.File`) - :type file_path: str :param destination: Filename, file path or instance of :class:`io.IOBase`. For e.g. :class:`io.BytesIO`, defaults to None - :type destination: Optional[Union[BinaryIO, pathlib.Path, str]] :param timeout: Total timeout in seconds, defaults to 30 - :type timeout: int :param chunk_size: File chunks size, defaults to 64 kb - :type chunk_size: int :param seek: Go to start of file when downloading is finished. Used only for destination with :class:`typing.BinaryIO` type, defaults to True - :type seek: bool """ if destination is None: destination = io.BytesIO() @@ -225,6 +233,41 @@ class Bot(ContextInstanceMixin["Bot"]): destination=destination, seek=seek, stream=stream ) + async def download( + self, + file: Union[str, Downloadable], + destination: Optional[Union[BinaryIO, pathlib.Path, str]] = None, + timeout: int = 30, + chunk_size: int = 65536, + seek: bool = True, + ) -> Optional[BinaryIO]: + """ + Download file by file_id or Downloadable object to destination. + + If you want to automatically create destination (:class:`io.BytesIO`) use default + value of destination and handle result of this method. + + :param file: file_id or Downloadable object + :param destination: Filename, file path or instance of :class:`io.IOBase`. For e.g. :class:`io.BytesIO`, defaults to None + :param timeout: Total timeout in seconds, defaults to 30 + :param chunk_size: File chunks size, defaults to 64 kb + :param seek: Go to start of file when downloading is finished. Used only for destination with :class:`typing.BinaryIO` type, defaults to True + """ + if isinstance(file, str): + file_id = file + elif hasattr(file, "file_id"): + file_id = file.file_id + else: + raise TypeError("file can only be of the string or Downloadable type") + + _file = await self.get_file(file_id) + # https://github.com/aiogram/aiogram/pull/282/files#r394110017 + file_path = cast(str, _file.file_path) + + return await self.download_file( + file_path, destination=destination, timeout=timeout, chunk_size=chunk_size, seek=seek + ) + async def __call__(self, method: TelegramMethod[T]) -> T: """ Call API method diff --git a/aiogram/api/types/downloadable.py b/aiogram/api/types/downloadable.py new file mode 100644 index 00000000..be808293 --- /dev/null +++ b/aiogram/api/types/downloadable.py @@ -0,0 +1,5 @@ +from typing import Protocol + + +class Downloadable(Protocol): + file_id: str diff --git a/docs/api/downloading_files.md b/docs/api/downloading_files.md index 21d6a3cc..7ec77600 100644 --- a/docs/api/downloading_files.md +++ b/docs/api/downloading_files.md @@ -13,11 +13,12 @@ file = await bot.get_file(file_id) file_path = file.file_path ``` -After that, use the `download_file` method from the bot object. +After that, use the [download_file](#download_file) method from the bot object. ### download_file(...) -Download file by file_path to destination. +Download file by `file_path` to destination. + If you want to automatically create destination (`#!python3 io.BytesIO`) use default value of destination and handle result of this method. @@ -54,4 +55,34 @@ If you leave the default value, an `#!python3 io.BytesIO` object will be created ```python3 result: io.BytesIO = await bot.download_file(file_path) -``` \ No newline at end of file +``` + +## Download file in short way + +Getting `file_path` manually every time is boring, so you should use the [download](#download) method. + +### download(...) + +Download file by `file_id` or `Downloadable` object to destination. + +If you want to automatically create destination (`#!python3 io.BytesIO`) use default +value of destination and handle result of this method. + +|Argument|Type|Description| +|---|---|---| +| file | `#!python3 Union[str, Downloadable]` | file_id or Downloadable object | +| destination | `#!python3 Optional[Union[BinaryIO, pathlib.Path, str]]` | Filename, file path or instance of `#!python3 io.IOBase`. For e.g. `#!python3 io.BytesIO` (Default: `#!python3 None`) | +| timeout | `#!python3 int` | Total timeout in seconds (Default: `30`) | +| chunk_size | `#!python3 int` | File chunks size (Default: `64 kb`) | +| seek | `#!python3 bool` | Go to start of file when downloading is finished. Used only for destination with `#!python3 typing.BinaryIO` type (Default: `#!python3 True`) | + +It differs from [download_file](#download_file) **only** in that it accepts `file_id` or an object that contains the `file_id` attribute instead of `file_path`. + +You can download a file to [disk](#download-file-to-disk) or to a [binary I/O](#download-file-to-binary-io-object) object in the same way. + +Example: + +```python3 +document = message.document +await bot.download(document) +``` diff --git a/tests/test_api/test_client/test_bot.py b/tests/test_api/test_client/test_bot.py index b5c59970..675c0dd3 100644 --- a/tests/test_api/test_client/test_bot.py +++ b/tests/test_api/test_client/test_bot.py @@ -6,7 +6,9 @@ from aresponses import ResponsesMockServer from aiogram import Bot from aiogram.api.client.session.aiohttp import AiohttpSession -from aiogram.api.methods import GetMe +from aiogram.api.methods import GetFile, GetMe +from aiogram.api.types import File, PhotoSize +from tests.mocked_bot import MockedBot try: from asynctest import CoroutineMock, patch @@ -112,3 +114,20 @@ class TestBot: assert isinstance(result, io.BytesIO) assert result is custom assert result.read() == b"\f" * 10 + + @pytest.mark.asyncio + async def test_download(self, bot: MockedBot, aresponses: ResponsesMockServer): + bot.add_result_for( + GetFile, ok=True, result=File(file_id="file id", file_unique_id="file id") + ) + bot.add_result_for( + GetFile, ok=True, result=File(file_id="file id", file_unique_id="file id") + ) + + assert await bot.download(File(file_id="file id", file_unique_id="file id")) + assert await bot.download("file id") + + with pytest.raises(TypeError): + await bot.download( + [PhotoSize(file_id="file id", file_unique_id="file id", width=123, height=123)] + )