Reworked InputFile sending (#1238)

* Reworked InputFile sending

* Added changelog
This commit is contained in:
Alex Root Junior 2023-08-02 21:41:07 +03:00 committed by GitHub
parent c7b7714959
commit a7916c1103
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 32 additions and 34 deletions

3
CHANGES/1238.misc.rst Normal file
View file

@ -0,0 +1,3 @@
Reworked InputFile reading, removed :code:`__aiter__` method, added `bot: Bot` argument to
the :code:`.read(...)` method, so, from now URLInputFile can be used without specifying
bot instance.

View file

@ -6,6 +6,7 @@ from typing import (
TYPE_CHECKING,
Any,
AsyncGenerator,
AsyncIterator,
Dict,
Iterable,
List,
@ -147,7 +148,11 @@ class AiohttpSession(BaseSession):
continue
form.add_field(key, value)
for key, value in files.items():
form.add_field(key, value, filename=value.filename or key)
form.add_field(
key,
value.read(bot),
filename=value.filename or key,
)
return form
async def make_request(

View file

@ -41,13 +41,9 @@ class InputFile(ABC):
self.chunk_size = chunk_size
@abstractmethod
async def read(self, chunk_size: int) -> AsyncGenerator[bytes, None]: # pragma: no cover
async def read(self, bot: "Bot") -> AsyncGenerator[bytes, None]: # pragma: no cover
yield b""
async def __aiter__(self) -> AsyncIterator[bytes]:
async for chunk in self.read(self.chunk_size):
yield chunk
class BufferedInputFile(InputFile):
def __init__(self, file: bytes, filename: str, chunk_size: int = DEFAULT_CHUNK_SIZE):
@ -84,9 +80,9 @@ class BufferedInputFile(InputFile):
data = f.read()
return cls(data, filename=filename, chunk_size=chunk_size)
async def read(self, chunk_size: int) -> AsyncGenerator[bytes, None]:
async def read(self, bot: "Bot") -> AsyncGenerator[bytes, None]:
buffer = io.BytesIO(self.data)
while chunk := buffer.read(chunk_size):
while chunk := buffer.read(self.chunk_size):
yield chunk
@ -111,9 +107,9 @@ class FSInputFile(InputFile):
self.path = path
async def read(self, chunk_size: int) -> AsyncGenerator[bytes, None]:
async def read(self, bot: "Bot") -> AsyncGenerator[bytes, None]:
async with aiofiles.open(self.path, "rb") as f:
while chunk := await f.read(chunk_size):
while chunk := await f.read(self.chunk_size):
yield chunk
@ -121,11 +117,11 @@ class URLInputFile(InputFile):
def __init__(
self,
url: str,
bot: "Bot",
headers: Optional[Dict[str, Any]] = None,
filename: Optional[str] = None,
chunk_size: int = DEFAULT_CHUNK_SIZE,
timeout: int = 30,
bot: Optional["Bot"] = None,
):
"""
Represents object for streaming files from internet
@ -136,7 +132,7 @@ class URLInputFile(InputFile):
:param chunk_size: Uploading chunk size
:param timeout: Timeout for downloading
:param bot: Bot instance to use HTTP session from.
If not specified, will be used current bot from context.
If not specified, will be used current bot
"""
super().__init__(filename=filename, chunk_size=chunk_size)
if headers is None:
@ -147,8 +143,9 @@ class URLInputFile(InputFile):
self.timeout = timeout
self.bot = bot
async def read(self, chunk_size: int) -> AsyncGenerator[bytes, None]:
stream = self.bot.session.stream_content(
async def read(self, bot: "Bot") -> AsyncGenerator[bytes, None]:
bot = self.bot or bot
stream = bot.session.stream_content(
url=self.url,
headers=self.headers,
timeout=self.timeout,

View file

@ -170,7 +170,7 @@ class BaseRequestHandler(ABC):
payload.set_content_disposition("form-data", name=key)
for key, value in files.items():
payload = writer.append(value)
payload = writer.append(value.read(bot))
payload.set_content_disposition(
"form-data",
name=key,

View file

@ -1,5 +1,5 @@
import asyncio
from typing import Any, AsyncContextManager, AsyncGenerator, Dict, List
from typing import Any, AsyncContextManager, AsyncGenerator, AsyncIterable, Dict, List
from unittest.mock import AsyncMock, patch
import aiohttp_socks
@ -156,7 +156,7 @@ class TestAiohttpSession:
assert fields[1][2].startswith("attach://")
assert fields[2][0]["name"] == fields[1][2][9:]
assert fields[2][0]["filename"] == "file.txt"
assert isinstance(fields[2][2], BareInputFile)
assert isinstance(fields[2][2], AsyncIterable)
async def test_make_request(self, bot: MockedBot, aresponses: ResponsesMockServer):
aresponses.add(

View file

@ -24,11 +24,7 @@ from aiogram.methods import (
SendVideoNote,
SendVoice,
)
from aiogram.types import (
Chat,
ChatJoinRequest,
User,
)
from aiogram.types import Chat, ChatJoinRequest, User
class TestChatJoinRequest:

View file

@ -12,19 +12,18 @@ class TestInputFile:
file = FSInputFile(__file__)
assert isinstance(file, InputFile)
assert isinstance(file, AsyncIterable)
assert file.filename is not None
assert file.filename.startswith("test_")
assert file.filename.endswith(".py")
assert file.chunk_size > 0
async def test_fs_input_file_readable(self):
async def test_fs_input_file_readable(self, bot: MockedBot):
file = FSInputFile(__file__, chunk_size=1)
assert file.chunk_size == 1
size = 0
async for chunk in file:
async for chunk in file.read(bot):
chunk_size = len(chunk)
assert chunk_size == 1
size += chunk_size
@ -34,15 +33,14 @@ class TestInputFile:
file = BufferedInputFile(b"\f" * 10, filename="file.bin")
assert isinstance(file, InputFile)
assert isinstance(file, AsyncIterable)
assert file.filename == "file.bin"
assert isinstance(file.data, bytes)
async def test_buffered_input_file_readable(self):
async def test_buffered_input_file_readable(self, bot: MockedBot):
file = BufferedInputFile(b"\f" * 10, filename="file.bin", chunk_size=1)
size = 0
async for chunk in file:
async for chunk in file.read(bot):
chunk_size = len(chunk)
assert chunk_size == 1
size += chunk_size
@ -52,32 +50,31 @@ class TestInputFile:
file = BufferedInputFile.from_file(__file__, chunk_size=10)
assert isinstance(file, InputFile)
assert isinstance(file, AsyncIterable)
assert file.filename is not None
assert file.filename.startswith("test_")
assert file.filename.endswith(".py")
assert isinstance(file.data, bytes)
assert file.chunk_size == 10
async def test_buffered_input_file_from_file_readable(self):
async def test_buffered_input_file_from_file_readable(self, bot: MockedBot):
file = BufferedInputFile.from_file(__file__, chunk_size=1)
size = 0
async for chunk in file:
async for chunk in file.read(bot):
chunk_size = len(chunk)
assert chunk_size == 1
size += chunk_size
assert size > 0
async def test_uri_input_file(self, aresponses: ResponsesMockServer):
async def test_url_input_file(self, aresponses: ResponsesMockServer):
aresponses.add(
aresponses.ANY, aresponses.ANY, "get", aresponses.Response(status=200, body=b"\f" * 10)
)
bot = Bot(token="42:TEST")
file = URLInputFile("https://test.org/", bot, chunk_size=1)
file = URLInputFile("https://test.org/", chunk_size=1)
size = 0
async for chunk in file:
async for chunk in file.read(bot):
assert chunk == b"\f"
chunk_size = len(chunk)
assert chunk_size == 1