diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index b39b11764..b34d0e893 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -212,8 +212,28 @@ def __init__(self, wrapped: io.IOBase, update: Callable[[int], None]) -> None: self._update = update def __iter__(self) -> Iterator[bytes]: + # For TextIOWrapper objects, use the underlying binary buffer position + # to track bytes consumed rather than character count. This matters for + # multi-byte encodings (e.g. utf-16-le) where len(line) is roughly half + # the actual byte count, causing the progress bar to stall at ~50%. + binary = getattr(self._wrapped, "buffer", None) + last_pos: Optional[int] = None + if binary is not None: + try: + last_pos = binary.tell() + except OSError: + binary = None for line in self._wrapped: - self._update(len(line)) + if binary is not None: + try: + pos = binary.tell() + self._update(pos - last_pos) + last_pos = pos + except OSError: + self._update(len(line)) + binary = None + else: + self._update(len(line)) yield line def read(self, size: int = -1) -> bytes: diff --git a/tests/test_cli.py b/tests/test_cli.py index a2135b087..b4f856349 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2205,6 +2205,42 @@ def test_insert_encoding(tmpdir): ] +def test_insert_utf16le_encoding(tmpdir): + """Progress bar stays accurate and data inserts correctly for UTF-16-LE CSV files. + + The progress bar was initialized with the file's byte length but UpdateWrapper + used len(decoded_line) (character count) for updates. For utf-16-le each + character is 2 bytes, so the bar stalled at ~50% instead of reaching 100%. + """ + db_path = str(tmpdir / "test.db") + csv_content = "id,name\n1,Alice\n2,Bob\n" + utf16le_bytes = csv_content.encode("utf-16-le") + csv_path = str(tmpdir / "test.csv") + with open(csv_path, "wb") as fp: + fp.write(utf16le_bytes) + + result = CliRunner().invoke( + cli.cli, + [ + "insert", + db_path, + "names", + csv_path, + "--csv", + "--encoding", + "utf-16-le", + "--no-detect-types", + ], + catch_exceptions=False, + ) + assert result.exit_code == 0 + db = Database(db_path) + assert list(db["names"].rows) == [ + {"id": "1", "name": "Alice"}, + {"id": "2", "name": "Bob"}, + ] + + @pytest.mark.parametrize("fts", ["FTS4", "FTS5"]) @pytest.mark.parametrize( "extra_arg,expected",