diff --git a/sqlite_utils/export.py b/sqlite_utils/export.py new file mode 100644 index 000000000..a5a8123f0 --- /dev/null +++ b/sqlite_utils/export.py @@ -0,0 +1,48 @@ +import csv +import itertools +from collections.abc import Mapping +from typing import Any, Sequence, TextIO, Type, Union + +CsvDialect = Union[str, csv.Dialect, Type[csv.Dialect]] + + +def rows_to_csv_file( + cursor: Any, + file: TextIO, + *, + header: bool = True, + dialect: CsvDialect = "excel", + **writer_kwargs: Any, +) -> None: + """Write the rows from a DB-API cursor to a CSV-compatible text stream. + + ``cursor`` should be a DB-API cursor whose ``description`` attribute + contains the result column metadata. Pass ``dialect="excel-tab"`` to + produce TSV output, or provide any dialect accepted by ``csv.writer``. + + Additional keyword arguments are forwarded to ``csv.writer``. + """ + description: Union[Sequence[Sequence[Any]], None] = cursor.description + if description is None: + raise ValueError("Cursor does not have result columns") + + columns = [column[0] for column in description] + row_iterator = iter(cursor) + sentinel = object() + first_row = next(row_iterator, sentinel) + if isinstance(first_row, Mapping) and len(set(columns)) != len(columns): + raise ValueError("Mapping rows cannot represent duplicate column names") + + writer = csv.writer(file, dialect=dialect, **writer_kwargs) + if header: + writer.writerow(columns) + + def normalize_row(row): + if isinstance(row, Mapping): + return [row[column] for column in columns] + return row + + if first_row is not sentinel: + writer.writerows( + normalize_row(row) for row in itertools.chain((first_row,), row_iterator) + ) diff --git a/tests/test_export.py b/tests/test_export.py new file mode 100644 index 000000000..44c08d4ed --- /dev/null +++ b/tests/test_export.py @@ -0,0 +1,101 @@ +import csv +import io +import sqlite3 + +import pytest + +from sqlite_utils.export import rows_to_csv_file + + +def _cursor(): + db = sqlite3.connect(":memory:") + db.execute("create table creatures (id integer, name text)") + db.executemany( + "insert into creatures (id, name) values (?, ?)", + ((1, "Cleo"), (2, "Cardi, Jr.")), + ) + return db, db.execute("select id, name from creatures order by id") + + +def test_rows_to_csv_file(): + db, cursor = _cursor() + output = io.StringIO(newline="") + rows_to_csv_file(cursor, output, lineterminator="\n") + assert output.getvalue() == 'id,name\n1,Cleo\n2,"Cardi, Jr."\n' + db.close() + + +def test_rows_to_csv_file_without_header(): + db, cursor = _cursor() + output = io.StringIO(newline="") + rows_to_csv_file(cursor, output, header=False, lineterminator="\n") + assert output.getvalue() == '1,Cleo\n2,"Cardi, Jr."\n' + db.close() + + +def test_rows_to_csv_file_tsv(): + db, cursor = _cursor() + output = io.StringIO(newline="") + rows_to_csv_file(cursor, output, dialect="excel-tab", lineterminator="\n") + assert output.getvalue() == "id\tname\n1\tCleo\n2\tCardi, Jr.\n" + db.close() + + +def test_rows_to_csv_file_accepts_dialect_class(): + class PipeDialect(csv.excel): + delimiter = "|" + + db, cursor = _cursor() + output = io.StringIO(newline="") + rows_to_csv_file(cursor, output, dialect=PipeDialect, lineterminator="\n") + assert output.getvalue() == "id|name\n1|Cleo\n2|Cardi, Jr.\n" + db.close() + + +def test_rows_to_csv_file_with_mapping_row_factory(): + db = sqlite3.connect(":memory:") + db.execute("create table creatures (id integer, name text)") + db.executemany( + "insert into creatures (id, name) values (?, ?)", + ((1, "Cleo"), (2, "Cardi, Jr.")), + ) + + def dict_factory(cursor, row): + return {column[0]: value for column, value in zip(cursor.description, row)} + + db.row_factory = dict_factory + cursor = db.execute("select id, name from creatures order by id") + output = io.StringIO(newline="") + rows_to_csv_file(cursor, output, lineterminator="\n") + assert output.getvalue() == 'id,name\n1,Cleo\n2,"Cardi, Jr."\n' + db.close() + + +def test_rows_to_csv_file_rejects_duplicate_mapping_columns(): + db = sqlite3.connect(":memory:") + db.execute("create table creatures (id integer, name text)") + db.execute("insert into creatures (id, name) values (1, 'Cleo')") + + def dict_factory(cursor, row): + return {column[0]: value for column, value in zip(cursor.description, row)} + + db.row_factory = dict_factory + cursor = db.execute("select id as value, name as value from creatures") + output = io.StringIO(newline="") + with pytest.raises( + ValueError, match="Mapping rows cannot represent duplicate column names" + ): + rows_to_csv_file(cursor, output, lineterminator="\n") + assert output.getvalue() == "" + db.close() + + +@pytest.mark.parametrize("header", (True, False)) +def test_rows_to_csv_file_requires_result_columns(header): + db = sqlite3.connect(":memory:") + cursor = db.execute("create table creatures (id integer)") + output = io.StringIO(newline="") + with pytest.raises(ValueError, match="Cursor does not have result columns"): + rows_to_csv_file(cursor, output, header=header) + assert output.getvalue() == "" + db.close()