diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index ee6695b55..4c8b3a46e 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -208,15 +208,39 @@ class UpdateWrapper: def __init__(self, wrapped: io.IOBase, update: Callable[[int], None]) -> None: self._wrapped = wrapped self._update = update + try: + self._position: int | None = wrapped.tell() + except (AttributeError, OSError): + self._position = None + + def _update_progress(self, fallback_length: int) -> None: + if self._position is None: + self._update(fallback_length) + return + try: + position = self._wrapped.tell() + except (AttributeError, OSError): + self._position = None + self._update(fallback_length) + return + delta = position - self._position + self._update(delta if delta >= 0 else fallback_length) + self._position = position def __iter__(self) -> Iterator[bytes]: - for line in self._wrapped: - self._update(len(line)) + # readline() keeps TextIOWrapper.tell() available, unlike iterating the + # TextIOWrapper directly. Its position is in the underlying file's + # bytes, so multibyte encodings advance the progress bar correctly. + while True: + line = self._wrapped.readline() + if not line: + break + self._update_progress(len(line)) yield line def read(self, size: int = -1) -> bytes: data = self._wrapped.read(size) - self._update(len(data)) + self._update_progress(len(data)) return data @@ -659,4 +683,4 @@ def flatten(row: dict[str, Any]) -> dict[str, Any]: :param row: A Python dictionary, optionally with nested dictionaries """ - return dict(_flatten(row)) + return dict(_flatten(row)) \ No newline at end of file diff --git a/tests/test_file_progress.py b/tests/test_file_progress.py new file mode 100644 index 000000000..4baf438f1 --- /dev/null +++ b/tests/test_file_progress.py @@ -0,0 +1,29 @@ +import io + +import pytest + +from sqlite_utils.utils import UpdateWrapper + + +@pytest.mark.parametrize("encoding", ("utf-8", "utf-16-le", "utf-32-le")) +def test_update_wrapper_tracks_underlying_bytes(encoding): + text = "id,name\r\n1,Café\r\n2,猫\r\n" + raw = text.encode(encoding) + updates = [] + decoded = io.TextIOWrapper(io.BytesIO(raw), encoding=encoding) + + lines = list(UpdateWrapper(decoded, updates.append)) + + assert lines == ["id,name\n", "1,Café\n", "2,猫\n"] + assert sum(updates) == len(raw) + + +def test_update_wrapper_read_tracks_underlying_bytes(): + raw = "é猫".encode("utf-8") + updates = [] + decoded = io.TextIOWrapper(io.BytesIO(raw), encoding="utf-8") + wrapped = UpdateWrapper(decoded, updates.append) + + assert wrapped.read(1) == "é" + assert wrapped.read() == "猫" + assert sum(updates) == len(raw)