Fix progress bar against non-UTF-8 encodings

For #439.

`file_progress()` sets the progress bar length to the file size in
bytes (`os.path.getsize(file.name)`), but `UpdateWrapper.__iter__`
called `update(len(line))`, which is the decoded character count.
With UTF-16-LE input every character is 2 bytes, so the bar capped
at 50%; UTF-32 capped at 25%; etc.

Simon noted in the issue that the obvious fix (calling `.tell()` on
the wrapped text stream) doesn't work because text mode disables it
during iteration. The underlying binary buffer doesn't have that
restriction though, so this tracks progress against
`TextIOWrapper.buffer.tell()` when the wrapped object exposes one.

For raw binary streams (no `.buffer` attribute) we keep the old
behaviour, which was already byte-accurate.

Added six regression tests in tests/test_utils.py covering UTF-8,
UTF-16-LE, BOM-prefixed UTF-16, the sniff-style `BufferedReader` chain,
a raw binary fallback, and the `.read()` path used by the JSON loader.
Each asserts that the sum of update() calls equals the on-disk file
size, which is what `click.progressbar` needs to reach 100%.
This commit is contained in:
LeSingh1 2026-05-18 16:29:36 -07:00
commit 625d004872
No known key found for this signature in database
2 changed files with 132 additions and 2 deletions

View file

@ -100,3 +100,96 @@ def test_flatten(input, expected):
)
def test_dedupe_keys(input, expected):
assert utils.dedupe_keys(input) == expected
# Regression tests for #439: progress bar against multi-byte encodings
def _collect_updates(rows):
"""Iterate the wrapper, capturing every update() value."""
return list(rows)
def _make_temp(content_bytes, tmp_path, name):
path = tmp_path / name
path.write_bytes(content_bytes)
return path
def test_updatewrapper_utf8_reports_byte_lengths(tmp_path):
# Sanity: ASCII / UTF-8 still hits 100% (this was already correct,
# but we want a baseline to protect.)
raw = b"a,b\n1,2\n3,4\n"
path = _make_temp(raw, tmp_path, "in.csv")
updates = []
with open(path, "rb") as fp:
wrapper = utils.UpdateWrapper(io.TextIOWrapper(fp, encoding="utf-8"), updates.append)
_collect_updates(wrapper)
assert sum(updates) == len(raw)
def test_updatewrapper_utf16le_reports_byte_lengths(tmp_path):
# Without the fix this test fails: the bar only reaches len(decoded)
# which is half the raw byte length for UTF-16-LE.
raw = "a,b\n1,2\n3,4\n".encode("utf-16-le")
path = _make_temp(raw, tmp_path, "in.csv")
updates = []
with open(path, "rb") as fp:
wrapper = utils.UpdateWrapper(io.TextIOWrapper(fp, encoding="utf-16-le"), updates.append)
_collect_updates(wrapper)
assert sum(updates) == len(raw)
def test_updatewrapper_utf16le_with_bom_reaches_total_bytes(tmp_path):
# BOM-prefixed UTF-16. The BOM byte is consumed by the TextIOWrapper
# before iteration starts; we should still account for the full file
# size so the bar reaches 100%.
raw = "" + "a,b\n1,2\n3,4\n"
raw_bytes = raw.encode("utf-16-le")
path = _make_temp(raw_bytes, tmp_path, "in.csv")
updates = []
with open(path, "rb") as fp:
wrapper = utils.UpdateWrapper(io.TextIOWrapper(fp, encoding="utf-16"), updates.append)
_collect_updates(wrapper)
assert sum(updates) == len(raw_bytes)
def test_updatewrapper_through_buffered_reader(tmp_path):
# The --sniff path wraps the raw file in io.BufferedReader before the
# TextIOWrapper. Progress reporting must still resolve to the binary
# file's byte count.
raw = "a,b\n1,2\n3,4\n".encode("utf-16-le")
path = _make_temp(raw, tmp_path, "in.csv")
updates = []
with open(path, "rb") as fp:
buffered = io.BufferedReader(fp, buffer_size=4096)
wrapper = utils.UpdateWrapper(
io.TextIOWrapper(buffered, encoding="utf-16-le"), updates.append
)
_collect_updates(wrapper)
assert sum(updates) == len(raw)
def test_updatewrapper_binary_file_unchanged(tmp_path):
# If the wrapped object is itself a raw binary file (no .buffer attr),
# we should keep the old behaviour: iterate yields bytes and len() is
# already the byte count.
raw = b"a,b\n1,2\n3,4\n"
path = _make_temp(raw, tmp_path, "in.csv")
updates = []
with open(path, "rb") as fp:
wrapper = utils.UpdateWrapper(fp, updates.append)
_collect_updates(wrapper)
assert sum(updates) == len(raw)
def test_updatewrapper_read_path_utf16le(tmp_path):
# The .read() path is used by the JSON loader (not the CSV iterator),
# but must agree with the iterator path on byte accounting.
raw = '{"a": 1}'.encode("utf-16-le")
path = _make_temp(raw, tmp_path, "in.json")
updates = []
with open(path, "rb") as fp:
wrapper = utils.UpdateWrapper(io.TextIOWrapper(fp, encoding="utf-16-le"), updates.append)
wrapper.read()
assert sum(updates) == len(raw)