mirror of
https://github.com/csunny/DB-GPT.git
synced 2026-07-16 17:15:22 +00:00
fix(core): StreamedBytesIO ignores negative read size and drops seek(SEEK_END) position (#3136)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -997,7 +997,7 @@ class StreamedBytesIO(io.BytesIO):
|
||||
bytes: The read data
|
||||
"""
|
||||
left_off_at = self._bytes.tell()
|
||||
if size is None:
|
||||
if size is None or size < 0:
|
||||
self._load_all()
|
||||
else:
|
||||
goal_position = left_off_at + size
|
||||
@@ -1006,20 +1006,22 @@ class StreamedBytesIO(io.BytesIO):
|
||||
self._bytes.seek(left_off_at)
|
||||
return self._bytes.read(size)
|
||||
|
||||
def seek(self, position: int, whence: int = io.SEEK_SET):
|
||||
def seek(self, position: int, whence: int = io.SEEK_SET) -> int:
|
||||
"""Seek to a position in the stream.
|
||||
|
||||
Args:
|
||||
position (int): The position
|
||||
whence (int, optional): The reference point. Defaults to io.SEEK
|
||||
|
||||
Returns:
|
||||
int: The new absolute position.
|
||||
|
||||
Raises:
|
||||
ValueError: If the reference point is invalid
|
||||
"""
|
||||
if whence == io.SEEK_END:
|
||||
self._load_all()
|
||||
else:
|
||||
self._bytes.seek(position, whence)
|
||||
return self._bytes.seek(position, whence)
|
||||
|
||||
def __enter__(self):
|
||||
"""Enter the context manager."""
|
||||
|
||||
@@ -13,6 +13,7 @@ from ..file import (
|
||||
InMemoryStorage,
|
||||
LocalFileStorage,
|
||||
SimpleDistributedStorage,
|
||||
StreamedBytesIO,
|
||||
)
|
||||
|
||||
|
||||
@@ -504,3 +505,16 @@ def test_simple_distributed_storage_delete_file_remote(
|
||||
f"http://{remote_node_address}/api/v2/serve/file/files/{bucket}/{file_id}",
|
||||
timeout=360,
|
||||
)
|
||||
|
||||
|
||||
class TestStreamedBytesIO:
|
||||
def test_read_negative_size_reads_all(self):
|
||||
stream = StreamedBytesIO(iter([b"hello ", b"world"]))
|
||||
assert stream.read(-1) == b"hello world"
|
||||
|
||||
def test_seek_end_loads_all_and_returns_position(self):
|
||||
stream = StreamedBytesIO(iter([b"hello ", b"world"]))
|
||||
position = stream.seek(-5, io.SEEK_END)
|
||||
assert position == 6
|
||||
assert stream.tell() == 6
|
||||
assert stream.read() == b"world"
|
||||
|
||||
Reference in New Issue
Block a user