diff --git a/sqlite_utils/export.py b/sqlite_utils/export.py index 9cde7dc..754fc74 100644 --- a/sqlite_utils/export.py +++ b/sqlite_utils/export.py @@ -1,7 +1,7 @@ import csv -from typing import Any, Sequence, TextIO, Union +from typing import Any, Sequence, TextIO, Type, Union -CsvDialect = Union[str, csv.Dialect] +CsvDialect = Union[str, csv.Dialect, Type[csv.Dialect]] def rows_to_csv_file( diff --git a/tests/test_export.py b/tests/test_export.py index a72b2bf..b2a0309 100644 --- a/tests/test_export.py +++ b/tests/test_export.py @@ -1,3 +1,4 @@ +import csv import io import sqlite3 @@ -40,6 +41,17 @@ def test_rows_to_csv_file_tsv(): 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() + + @pytest.mark.parametrize("header", (True, False)) def test_rows_to_csv_file_requires_result_columns(header): db = sqlite3.connect(":memory:")