mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-07-23 09:24:31 +02:00
9a605e058cUpdateWrapper tracked progress using len(line) which counts characters, but the progress bar total is set to the file size in bytes. For multi-byte encodings like UTF-16-LE (2 bytes per character) the bar would only reach 50% before the file finished. Now encode each line/chunk back to bytes using the stream's encoding to report the correct byte count.
135 lines
4.2 KiB
Python
135 lines
4.2 KiB
Python
from sqlite_utils import utils
|
|
import csv
|
|
import io
|
|
import pytest
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input,expected,should_be_is",
|
|
[
|
|
({}, None, True),
|
|
({"foo": "bar"}, None, True),
|
|
(
|
|
{"content": {"$base64": True, "encoded": "aGVsbG8="}},
|
|
{"content": b"hello"},
|
|
False,
|
|
),
|
|
],
|
|
)
|
|
def test_decode_base64_values(input, expected, should_be_is):
|
|
actual = utils.decode_base64_values(input)
|
|
if should_be_is:
|
|
assert actual is input
|
|
else:
|
|
assert actual == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"size,expected",
|
|
(
|
|
(1, [["a"], ["b"], ["c"], ["d"]]),
|
|
(2, [["a", "b"], ["c", "d"]]),
|
|
(3, [["a", "b", "c"], ["d"]]),
|
|
(4, [["a", "b", "c", "d"]]),
|
|
),
|
|
)
|
|
def test_chunks(size, expected):
|
|
input = ["a", "b", "c", "d"]
|
|
chunks = list(map(list, utils.chunks(input, size)))
|
|
assert chunks == expected
|
|
|
|
|
|
def test_hash_record():
|
|
expected = "d383e7c0ba88f5ffcdd09be660de164b3847401a"
|
|
assert utils.hash_record({"name": "Cleo", "twitter": "CleoPaws"}) == expected
|
|
assert (
|
|
utils.hash_record(
|
|
{"name": "Cleo", "twitter": "CleoPaws", "age": 7}, keys=("name", "twitter")
|
|
)
|
|
== expected
|
|
)
|
|
assert (
|
|
utils.hash_record({"name": "Cleo", "twitter": "CleoPaws", "age": 7}) != expected
|
|
)
|
|
|
|
|
|
def test_maximize_csv_field_size_limit():
|
|
# Reset to default in case other tests have changed it
|
|
csv.field_size_limit(utils.ORIGINAL_CSV_FIELD_SIZE_LIMIT)
|
|
long_value = "a" * 131073
|
|
long_csv = "id,text\n1,{}".format(long_value)
|
|
fp = io.BytesIO(long_csv.encode("utf-8"))
|
|
# Using rows_from_file should error
|
|
with pytest.raises(csv.Error):
|
|
rows, _ = utils.rows_from_file(fp, utils.Format.CSV)
|
|
list(rows)
|
|
# But if we call maximize_csv_field_size_limit() first it should be OK:
|
|
utils.maximize_csv_field_size_limit()
|
|
fp2 = io.BytesIO(long_csv.encode("utf-8"))
|
|
rows2, _ = utils.rows_from_file(fp2, utils.Format.CSV)
|
|
rows_list2 = list(rows2)
|
|
assert len(rows_list2) == 1
|
|
assert rows_list2[0]["id"] == "1"
|
|
assert rows_list2[0]["text"] == long_value
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input,expected",
|
|
(
|
|
({"foo": {"bar": 1}}, {"foo_bar": 1}),
|
|
({"foo": {"bar": [1, 2, {"baz": 3}]}}, {"foo_bar": [1, 2, {"baz": 3}]}),
|
|
({"foo": {"bar": 1, "baz": {"three": 3}}}, {"foo_bar": 1, "foo_baz_three": 3}),
|
|
),
|
|
)
|
|
def test_flatten(input, expected):
|
|
assert utils.flatten(input) == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input,expected",
|
|
(
|
|
([], []),
|
|
(["id", "name"], ["id", "name"]),
|
|
(["id", "id"], ["id", "id_2"]),
|
|
(["id", "id", "id"], ["id", "id_2", "id_3"]),
|
|
# A renamed duplicate must not clobber a real column called id_2
|
|
(["id", "id", "id_2"], ["id", "id_3", "id_2"]),
|
|
(["id_2", "id", "id"], ["id_2", "id", "id_3"]),
|
|
(["id", "id", "id_2", "id_2"], ["id", "id_3", "id_2", "id_2_2"]),
|
|
),
|
|
)
|
|
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)
|