Fix readable BufferedInputFile and add support BufferedIOBase

This commit is contained in:
teri-anric 2024-08-17 23:25:32 +03:00
parent f10414c1eb
commit 0ae2d82451
3 changed files with 41 additions and 9 deletions

5
CHANGES/1564.misc.rst Normal file
View 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.

View file

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

View file

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