This commit is contained in:
Chen Yufeiyang 2026-10-11 11:36:20 +08:00 • committed by GitHub
commit d9b281c712
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 50 additions and 18 deletions

View file

@ -358,17 +358,18 @@ def rows_from_file(
)
return _CloseableIterator(iter(rows), decoded_fp), Format.CSV
elif format == Format.TSV:
# The inner call applies the extra-field strategy, so the
# ignore_extras= and extras_key= arguments must be passed to it -
# see https://github.com/simonw/sqlite-utils/issues/892
rows, _ = rows_from_file(
fp, format=Format.CSV, dialect=csv.excel_tab, encoding=encoding
)
return (
_extra_key_strategy(
cast(Iterable[dict[str | None, object]], rows),
ignore_extras,
extras_key,
),
Format.TSV,
fp,
format=Format.CSV,
dialect=csv.excel_tab,
encoding=encoding,
ignore_extras=ignore_extras,
extras_key=extras_key,
)
return rows, Format.TSV
elif format is None:
# Detect the format, then call this recursively
buffered = io.BufferedReader(cast(io.RawIOBase, fp), buffer_size=4096)
@ -393,18 +394,16 @@ def rows_from_file(
first_bytes.decode(encoding or "utf-8-sig", "ignore")
)
rows, _ = rows_from_file(
buffered, format=Format.CSV, dialect=dialect, encoding=encoding
buffered,
format=Format.CSV,
dialect=dialect,
encoding=encoding,
ignore_extras=ignore_extras,
extras_key=extras_key,
)
# Make sure we return the format we detected
detected_format = Format.TSV if dialect.delimiter == "\t" else Format.CSV
return (
_extra_key_strategy(
cast(Iterable[dict[str | None, object]], rows),
ignore_extras,
extras_key,
),
detected_format,
)
return rows, detected_format
else:
raise RowsFromFileError("Bad format")

View file

@ -57,6 +57,39 @@ def test_rows_from_file_extra_fields_strategies(ignore_extras, extras_key, expec
assert list_rows == expected
@pytest.mark.parametrize(
"ignore_extras,extras_key,expected",
(
(True, None, [{"id": "1", "name": "Cleo"}]),
(False, "_rest", [{"id": "1", "name": "Cleo", "_rest": ["oops"]}]),
# expected of None means expect an error:
(False, False, None),
),
)
def test_rows_from_file_tsv_extra_fields_strategies(
ignore_extras, extras_key, expected
):
# ignore_extras= and extras_key= must apply to TSV as well as CSV,
# see https://github.com/simonw/sqlite-utils/issues/892
try:
rows, detected_format = rows_from_file(
BytesIO(b"id\tname\r\n1\tCleo\toops"),
format=Format.TSV,
ignore_extras=ignore_extras,
extras_key=extras_key,
)
list_rows = list(rows)
except RowError:
if expected is None:
# This is fine,
return
else:
# We did not expect an error
raise
assert detected_format == Format.TSV
assert list_rows == expected
def test_rows_from_file_error_on_string_io():
with pytest.raises(TypeError) as ex:
rows_from_file(StringIO("id,name\r\n1,Cleo")) # type: ignore[arg-type]