diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index b39b11764..3135c27af 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -212,13 +212,15 @@ def __init__(self, wrapped: io.IOBase, update: Callable[[int], None]) -> None: self._update = update def __iter__(self) -> Iterator[bytes]: + encoding = getattr(self._wrapped, "encoding", None) or "utf-8" for line in self._wrapped: - self._update(len(line)) + self._update(len(line.encode(encoding, errors="replace"))) yield line def read(self, size: int = -1) -> bytes: data = self._wrapped.read(size) - self._update(len(data)) + encoding = getattr(self._wrapped, "encoding", None) or "utf-8" + self._update(len(data.encode(encoding, errors="replace"))) return data diff --git a/tests/test_utils.py b/tests/test_utils.py index 3de5e94a8..653b0e448 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -100,3 +100,36 @@ def test_flatten(input, expected): ) def test_dedupe_keys(input, expected): assert utils.dedupe_keys(input) == expected + + +def test_update_wrapper_counts_bytes_not_chars(): + # For utf-16-le each ASCII character is 2 bytes, so character count != byte count. + # UpdateWrapper must report byte counts to match the os.path.getsize() total used by + # the progress bar; otherwise the bar stops at ~50% for utf-16-le files. + lines = ["hello\n", "world\n"] + + class FakeTextFile: + encoding = "utf-16-le" + + def __iter__(self): + return iter(lines) + + def read(self, size=-1): + return "".join(lines) + + byte_counts = [] + wrapper = utils.UpdateWrapper(FakeTextFile(), byte_counts.append) + + # Test __iter__ + assert list(wrapper) == lines + expected_iter_counts = [len(line.encode("utf-16-le")) for line in lines] + assert byte_counts == expected_iter_counts + # Sanity check: byte counts must differ from character counts for utf-16-le + assert byte_counts[0] == len("hello\n") * 2 + + # Test read() + byte_counts.clear() + wrapper2 = utils.UpdateWrapper(FakeTextFile(), byte_counts.append) + data = wrapper2.read() + assert data == "hello\nworld\n" + assert byte_counts == [len("hello\nworld\n".encode("utf-16-le"))]