diff --git a/sqlite_utils/export.py b/sqlite_utils/export.py index cf913d4..44dbcad 100644 --- a/sqlite_utils/export.py +++ b/sqlite_utils/export.py @@ -1,5 +1,5 @@ import csv -from typing import Any, Iterable, Optional, Sequence, TextIO, Union +from typing import Any, Sequence, TextIO, Union CsvDialect = Union[str, csv.Dialect] @@ -21,16 +21,11 @@ def rows_to_csv_file( 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") + writer = csv.writer(file, dialect=dialect, **writer_kwargs) if header: - description: Optional[Sequence[Sequence[Any]]] = cursor.description - if description is None: - raise ValueError("Cursor does not have result columns") writer.writerow([column[0] for column in description]) - writer.writerows(_rows(cursor)) - - -def _rows(cursor: Iterable[Sequence[Any]]) -> Iterable[Sequence[Any]]: - """Keep row iteration lazy so large query results are streamed.""" - for row in cursor: - yield row + writer.writerows(cursor) diff --git a/tests/test_export.py b/tests/test_export.py index bb91022..a72b2bf 100644 --- a/tests/test_export.py +++ b/tests/test_export.py @@ -40,10 +40,12 @@ def test_rows_to_csv_file_tsv(): db.close() -def test_rows_to_csv_file_requires_result_columns_for_header(): +@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) + rows_to_csv_file(cursor, output, header=header) + assert output.getvalue() == "" db.close()