mirror of
https://github.com/aiogram/aiogram.git
synced 2026-04-08 16:37:47 +00:00
Fix readable BufferedInputFile and add support BufferedIOBase
This commit is contained in:
parent
f10414c1eb
commit
0ae2d82451
3 changed files with 41 additions and 9 deletions
5
CHANGES/1564.misc.rst
Normal file
5
CHANGES/1564.misc.rst
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
Fixed unsafe file reading in BufferedInputFile.from_file and
|
||||
added support for initializing BufferedInputFile from BytesIO and
|
||||
other bytes buffers.
|
||||
This change improves safety and extends flexibility in handling
|
||||
different byte sources when working with buffered input files.
|
||||
|
|
@ -38,7 +38,12 @@ class InputFile(ABC):
|
|||
|
||||
|
||||
class BufferedInputFile(InputFile):
|
||||
def __init__(self, file: bytes, filename: str, chunk_size: int = DEFAULT_CHUNK_SIZE):
|
||||
def __init__(
|
||||
self,
|
||||
buffer_or_bytes: bytes | io.BufferedIOBase,
|
||||
filename: str,
|
||||
chunk_size: int = DEFAULT_CHUNK_SIZE,
|
||||
):
|
||||
"""
|
||||
Represents object for uploading files from filesystem
|
||||
|
||||
|
|
@ -48,7 +53,10 @@ class BufferedInputFile(InputFile):
|
|||
"""
|
||||
super().__init__(filename=filename, chunk_size=chunk_size)
|
||||
|
||||
self.data = file
|
||||
if isinstance(buffer_or_bytes, bytes):
|
||||
buffer_or_bytes = io.BytesIO(buffer_or_bytes)
|
||||
|
||||
self.buffer = buffer_or_bytes
|
||||
|
||||
@classmethod
|
||||
def from_file(
|
||||
|
|
@ -68,13 +76,15 @@ class BufferedInputFile(InputFile):
|
|||
"""
|
||||
if filename is None:
|
||||
filename = os.path.basename(path)
|
||||
with open(path, "rb") as f:
|
||||
data = f.read()
|
||||
return cls(data, filename=filename, chunk_size=chunk_size)
|
||||
buffer = io.BytesIO()
|
||||
with open(path, "rb") as fp:
|
||||
while chunk := fp.read(chunk_size):
|
||||
buffer.write(chunk)
|
||||
return cls(buffer, filename=filename, chunk_size=chunk_size)
|
||||
|
||||
async def read(self, bot: "Bot") -> AsyncGenerator[bytes, None]:
|
||||
buffer = io.BytesIO(self.data)
|
||||
while chunk := buffer.read(self.chunk_size):
|
||||
self.buffer.seek(0)
|
||||
while chunk := self.buffer.read(self.chunk_size):
|
||||
yield chunk
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import io
|
||||
from typing import AsyncIterable
|
||||
|
||||
from aresponses import ResponsesMockServer
|
||||
|
|
@ -34,7 +35,7 @@ class TestInputFile:
|
|||
|
||||
assert isinstance(file, InputFile)
|
||||
assert file.filename == "file.bin"
|
||||
assert isinstance(file.data, bytes)
|
||||
assert isinstance(file.buffer, io.BytesIO)
|
||||
|
||||
async def test_buffered_input_file_readable(self, bot: MockedBot):
|
||||
file = BufferedInputFile(b"\f" * 10, filename="file.bin", chunk_size=1)
|
||||
|
|
@ -53,7 +54,7 @@ class TestInputFile:
|
|||
assert file.filename is not None
|
||||
assert file.filename.startswith("test_")
|
||||
assert file.filename.endswith(".py")
|
||||
assert isinstance(file.data, bytes)
|
||||
assert isinstance(file.buffer, io.BytesIO)
|
||||
assert file.chunk_size == 10
|
||||
|
||||
async def test_buffered_input_file_from_file_readable(self, bot: MockedBot):
|
||||
|
|
@ -66,6 +67,22 @@ class TestInputFile:
|
|||
size += chunk_size
|
||||
assert size > 0
|
||||
|
||||
async def test_buffered_input_file_from_buffer(self, bot: MockedBot):
|
||||
buffer = io.BytesIO(b"\f" * 10)
|
||||
file = BufferedInputFile(buffer, filename="file.bin", chunk_size=1)
|
||||
|
||||
assert file.buffer is buffer
|
||||
assert file.filename == "file.bin"
|
||||
|
||||
size = 0
|
||||
async for chunk in file.read(bot):
|
||||
chunk_size = len(chunk)
|
||||
assert isinstance(chunk, bytes)
|
||||
assert chunk_size == 1
|
||||
size += chunk_size
|
||||
assert size == 10
|
||||
assert file.buffer.seek(0, io.SEEK_CUR) == 10
|
||||
|
||||
async def test_url_input_file(self, aresponses: ResponsesMockServer):
|
||||
aresponses.add(
|
||||
aresponses.ANY,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue