mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-09-08 01:24:20 +02:00
74c7038b8b
The CLI can render query/table output as CSV via --csv, but the Python API had no equivalent. sqlite_utils.utils.rows_to_csv() takes rows from db.query() or table.rows/rows_where() and writes CSV to a file-like object, or returns it as a string, with a header row by default (skip with no_headers=True) to mirror the CLI's behavior.
126 lines
3.7 KiB
Python
126 lines
3.7 KiB
Python
import csv
|
|
import io
|
|
|
|
import pytest
|
|
|
|
from sqlite_utils import utils
|
|
|
|
|
|
@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 = f"id,text\n1,{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_rows_to_csv_returns_string():
|
|
rows = [{"id": 1, "name": "Cleo"}, {"id": 2, "name": "Pancakes"}]
|
|
assert utils.rows_to_csv(rows) == "id,name\r\n1,Cleo\r\n2,Pancakes\r\n"
|
|
|
|
|
|
def test_rows_to_csv_writes_to_fp():
|
|
fp = io.StringIO()
|
|
result = utils.rows_to_csv([{"id": 1, "name": "Cleo"}], fp=fp)
|
|
assert result is None
|
|
assert fp.getvalue() == "id,name\r\n1,Cleo\r\n"
|
|
|
|
|
|
def test_rows_to_csv_no_headers():
|
|
csv_string = utils.rows_to_csv([{"id": 1, "name": "Cleo"}], no_headers=True)
|
|
assert csv_string == "1,Cleo\r\n"
|
|
|
|
|
|
def test_rows_to_csv_empty_rows():
|
|
assert utils.rows_to_csv([]) == ""
|
|
assert utils.rows_to_csv([], headers=["id", "name"]) == "id,name\r\n"
|