Add Python CSV export helper

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)
+        )

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()