mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-09-28 12:54:15 +02:00
Add rows_to_csv() Python API helper for writing rows as CSV (issue #580)
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.
This commit is contained in:
parent
71664b8fcf
commit
74c7038b8b
4 changed files with 105 additions and 0 deletions
|
|
@ -3307,6 +3307,29 @@ The ``sqlite_utils.utils.rows_from_file()`` helper function can read rows (a seq
|
||||||
.. autofunction:: sqlite_utils.utils.rows_from_file
|
.. autofunction:: sqlite_utils.utils.rows_from_file
|
||||||
:noindex:
|
:noindex:
|
||||||
|
|
||||||
|
.. _python_api_rows_to_csv:
|
||||||
|
|
||||||
|
Writing rows to CSV
|
||||||
|
===================
|
||||||
|
|
||||||
|
The CLI can render query and table output as CSV using ``--csv``, see :ref:`cli_query_csv`. The ``sqlite_utils.utils.rows_to_csv()`` helper provides the same behavior from the Python API - pass it rows from :meth:`db.query() <python_api_query>` or ``table.rows`` / ``table.rows_where()`` and it writes CSV to a file-like object, or returns the CSV as a string if you don't pass one:
|
||||||
|
|
||||||
|
.. code-block:: python
|
||||||
|
|
||||||
|
from sqlite_utils.utils import rows_to_csv
|
||||||
|
|
||||||
|
db = Database("dogs.db")
|
||||||
|
csv_string = rows_to_csv(db["dogs"].rows)
|
||||||
|
|
||||||
|
# Or write directly to a file:
|
||||||
|
with open("dogs.csv", "w", newline="") as fp:
|
||||||
|
rows_to_csv(db["dogs"].rows, fp=fp)
|
||||||
|
|
||||||
|
Like the CLI, a header row is included by default - taken from the keys of the first row - and can be skipped with ``no_headers=True``.
|
||||||
|
|
||||||
|
.. autofunction:: sqlite_utils.utils.rows_to_csv
|
||||||
|
:noindex:
|
||||||
|
|
||||||
.. _python_api_maximize_csv_field_size_limit:
|
.. _python_api_maximize_csv_field_size_limit:
|
||||||
|
|
||||||
Setting the maximum CSV field size limit
|
Setting the maximum CSV field size limit
|
||||||
|
|
|
||||||
|
|
@ -94,6 +94,13 @@ sqlite_utils.utils.rows_from_file
|
||||||
|
|
||||||
.. autofunction:: sqlite_utils.utils.rows_from_file
|
.. autofunction:: sqlite_utils.utils.rows_from_file
|
||||||
|
|
||||||
|
.. _reference_utils_rows_to_csv:
|
||||||
|
|
||||||
|
sqlite_utils.utils.rows_to_csv
|
||||||
|
------------------------------
|
||||||
|
|
||||||
|
.. autofunction:: sqlite_utils.utils.rows_to_csv
|
||||||
|
|
||||||
.. _reference_utils_typetracker:
|
.. _reference_utils_typetracker:
|
||||||
|
|
||||||
sqlite_utils.utils.TypeTracker
|
sqlite_utils.utils.TypeTracker
|
||||||
|
|
|
||||||
|
|
@ -14,6 +14,7 @@ from typing import (
|
||||||
TYPE_CHECKING,
|
TYPE_CHECKING,
|
||||||
Any,
|
Any,
|
||||||
BinaryIO,
|
BinaryIO,
|
||||||
|
TextIO,
|
||||||
TypeVar,
|
TypeVar,
|
||||||
Union,
|
Union,
|
||||||
cast,
|
cast,
|
||||||
|
|
@ -392,6 +393,58 @@ def rows_from_file(
|
||||||
raise RowsFromFileError("Bad format")
|
raise RowsFromFileError("Bad format")
|
||||||
|
|
||||||
|
|
||||||
|
def rows_to_csv(
|
||||||
|
rows: Iterable[Row],
|
||||||
|
fp: TextIO | None = None,
|
||||||
|
headers: Iterable[str] | None = None,
|
||||||
|
dialect: str | type[csv.Dialect] = "excel",
|
||||||
|
no_headers: bool = False,
|
||||||
|
) -> str | None:
|
||||||
|
"""
|
||||||
|
Write a sequence of dictionaries - such as rows from :meth:`.Database.query`
|
||||||
|
or :attr:`.Table.rows` - to CSV, mirroring the CLI's ``--csv`` output.
|
||||||
|
|
||||||
|
.. code-block:: python
|
||||||
|
|
||||||
|
from sqlite_utils.utils import rows_to_csv
|
||||||
|
|
||||||
|
csv_string = rows_to_csv([{"id": 1, "name": "Cleo"}])
|
||||||
|
print(csv_string)
|
||||||
|
# Outputs "id,name\\r\\n1,Cleo\\r\\n"
|
||||||
|
|
||||||
|
Pass ``fp=`` a writable file-like object to write there instead - in that
|
||||||
|
case this function returns ``None``.
|
||||||
|
|
||||||
|
:param rows: iterable of dictionaries to write
|
||||||
|
:param fp: optional writable file-like object - if omitted the CSV is
|
||||||
|
returned as a string
|
||||||
|
:param headers: explicit list of column headers - defaults to the keys
|
||||||
|
of the first row, the same as the CLI
|
||||||
|
:param dialect: the CSV dialect to use, defaults to ``"excel"``
|
||||||
|
:param no_headers: set to ``True`` to skip the header row, equivalent to
|
||||||
|
the CLI's ``--no-headers`` option
|
||||||
|
"""
|
||||||
|
return_string = fp is None
|
||||||
|
out: TextIO = fp if fp is not None else io.StringIO()
|
||||||
|
rows_iter = iter(rows)
|
||||||
|
first_row = next(rows_iter, None)
|
||||||
|
if headers is None:
|
||||||
|
if first_row is None:
|
||||||
|
# Nothing to write and no columns were specified
|
||||||
|
return cast(io.StringIO, out).getvalue() if return_string else None
|
||||||
|
headers = list(first_row.keys())
|
||||||
|
writer = csv.DictWriter(
|
||||||
|
out, fieldnames=list(headers), dialect=dialect, extrasaction="ignore"
|
||||||
|
)
|
||||||
|
if not no_headers:
|
||||||
|
writer.writeheader()
|
||||||
|
if first_row is not None:
|
||||||
|
writer.writerow(first_row)
|
||||||
|
for row in rows_iter:
|
||||||
|
writer.writerow(row)
|
||||||
|
return cast(io.StringIO, out).getvalue() if return_string else None
|
||||||
|
|
||||||
|
|
||||||
class TypeTracker:
|
class TypeTracker:
|
||||||
"""
|
"""
|
||||||
Wrap an iterator of dictionaries and keep track of which SQLite column
|
Wrap an iterator of dictionaries and keep track of which SQLite column
|
||||||
|
|
|
||||||
|
|
@ -102,3 +102,25 @@ def test_flatten(input, expected):
|
||||||
)
|
)
|
||||||
def test_dedupe_keys(input, expected):
|
def test_dedupe_keys(input, expected):
|
||||||
assert utils.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"
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue