Merge 2402094e5e into 9c66bafa52

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
 
 
@@ -663,4 +687,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))

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)