Merge 9a605e058c into a947dc6739

sqlite_utils/utils.py

@@ -210,15 +210,23 @@ class UpdateWrapper:
     def __init__(self, wrapped: io.IOBase, update: Callable[[int], None]) -> None:
         self._wrapped = wrapped
         self._update = update
+        # If this wraps a text stream, encode lines back to bytes for accurate
+        # byte-level progress tracking (fixes #439 for multi-byte encodings).
+        self._encoding: Optional[str] = getattr(wrapped, "encoding", None)
 
-    def __iter__(self) -> Iterator[bytes]:
+    def _byte_len(self, s: Any) -> int:
+        if self._encoding and isinstance(s, str):
+            return len(s.encode(self._encoding, errors="replace"))
+        return len(s)
+
+    def __iter__(self) -> Iterator[Any]:
         for line in self._wrapped:
-            self._update(len(line))
+            self._update(self._byte_len(line))
             yield line
 
-    def read(self, size: int = -1) -> bytes:
+    def read(self, size: int = -1) -> Any:
         data = self._wrapped.read(size)
-        self._update(len(data))
+        self._update(self._byte_len(data))
         return data
 
 

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_byte_progress_for_multibyte_encoding():
+    """UpdateWrapper should track bytes consumed, not characters (issue #439)."""
+    content = "hello\nworld\n"
+    # Build a UTF-16-LE encoded binary stream
+    encoded = content.encode("utf-16-le")
+    binary_stream = io.BytesIO(encoded)
+    text_stream = io.TextIOWrapper(binary_stream, encoding="utf-16-le")
+
+    updates = []
+    wrapper = utils.UpdateWrapper(text_stream, updates.append)
+
+    lines = list(wrapper)
+    assert lines == ["hello\n", "world\n"]
+
+    # Each character is 2 bytes in UTF-16-LE, so total bytes == 2 * len(content)
+    assert sum(updates) == len(encoded)
+
+
+def test_update_wrapper_byte_progress_for_utf8():
+    """UpdateWrapper byte tracking should also be correct for UTF-8."""
+    content = "hello\nworld\n"
+    encoded = content.encode("utf-8")
+    binary_stream = io.BytesIO(encoded)
+    text_stream = io.TextIOWrapper(binary_stream, encoding="utf-8")
+
+    updates = []
+    wrapper = utils.UpdateWrapper(text_stream, updates.append)
+
+    lines = list(wrapper)
+    assert lines == ["hello\n", "world\n"]
+    assert sum(updates) == len(encoded)