diff --git a/docs/examples/export_entire_dataset_to_file.mdx b/docs/examples/export_entire_dataset_to_file.mdx index 5cf4a2da77..17c2a10a68 100644 --- a/docs/examples/export_entire_dataset_to_file.mdx +++ b/docs/examples/export_entire_dataset_to_file.mdx @@ -11,7 +11,7 @@ import RunnableCodeBlock from '@site/src/components/RunnableCodeBlock'; import JsonExample from '!!raw-loader!roa-loader!./code_examples/export_entire_dataset_to_file_json.py'; import CsvExample from '!!raw-loader!roa-loader!./code_examples/export_entire_dataset_to_file_csv.py'; -This example demonstrates how to use the `BasicCrawler.export_data` method of the crawler to export the entire default dataset to a single file. This method supports exporting data in either CSV or JSON format and also accepts additional keyword arguments so you can fine-tune the underlying `json.dump` or `csv.writer` behavior. +This example demonstrates how to use the `BasicCrawler.export_data` method of the crawler to export the entire default dataset to a single file. This method supports exporting data in either CSV or JSON format and also accepts additional keyword arguments so you can fine-tune the underlying `json.dump` or `csv.DictWriter` behavior. :::note @@ -31,3 +31,17 @@ For these examples, we are using the `B + +## CSV columns + +Dataset items don't have to share a schema. Different handlers can push different fields, so the items of one dataset often have different keys. + +By default, the CSV columns are the keys of the first non-empty item. When a later item has a key the first one doesn't, its value isn't written, and Crawlee logs a warning naming the dropped keys. To use the keys of all items as columns instead, pass `collect_all_keys=True`: + +```python +await crawler.export_data(path='results.csv', collect_all_keys=True) +``` + +No value is dropped in that mode. Collecting all keys means reading the whole dataset before the first row can be written, so only the default mode writes rows as it goes. + +In both modes, cells for columns an item doesn't have stay empty, or hold the `restval` value if you pass one. diff --git a/src/crawlee/_types.py b/src/crawlee/_types.py index ffeb92094c..6bd500e3a2 100644 --- a/src/crawlee/_types.py +++ b/src/crawlee/_types.py @@ -772,8 +772,12 @@ class ExportDataJsonKwargs(TypedDict): """Specifies whether the output JSON object should have keys sorted alphabetically.""" -class ExportDataCsvKwargs(TypedDict): - """Keyword arguments for dataset's `export_data_csv` method.""" +class ExportDataCsvWriterKwargs(TypedDict): + """Keyword arguments forwarded to the underlying `csv` writer. + + The `fieldnames` and `extrasaction` arguments of `csv.DictWriter` are managed by Crawlee and cannot be + overridden. Use `collect_all_keys` to control which columns the export contains. + """ dialect: NotRequired[str] """Specifies a dialect to be used in CSV parsing and writing.""" @@ -800,6 +804,9 @@ class ExportDataCsvKwargs(TypedDict): """Controls when quotes should be generated by the writer and recognized by the reader. Can take any of the `QUOTE_*` constants, with a default of `QUOTE_MINIMAL`.""" + restval: NotRequired[str] + """The value written for columns an item does not contain. Defaults to an empty string.""" + skipinitialspace: NotRequired[bool] """When True, spaces immediately following the delimiter are ignored. Defaults to False.""" @@ -807,6 +814,15 @@ class ExportDataCsvKwargs(TypedDict): """When True, raises an exception on bad CSV input. Defaults to False.""" +class ExportDataCsvKwargs(ExportDataCsvWriterKwargs): + """Keyword arguments for dataset's CSV export.""" + + collect_all_keys: NotRequired[bool] + """When True, the columns are the keys of all items combined, and no value is dropped. When False, the columns + are the keys of the first non-empty item, and keys introduced by later items are dropped with a warning. + Defaults to False.""" + + class ExportDataKwargs(ExportDataJsonKwargs, ExportDataCsvKwargs): """Keyword arguments accepted by `BasicCrawler.export_data`. diff --git a/src/crawlee/_utils/file.py b/src/crawlee/_utils/file.py index 499d363046..9d2f9177cc 100644 --- a/src/crawlee/_utils/file.py +++ b/src/crawlee/_utils/file.py @@ -6,8 +6,9 @@ import os import sys import tempfile +from logging import getLogger from pathlib import Path -from typing import TYPE_CHECKING, overload +from typing import TYPE_CHECKING, cast, overload if TYPE_CHECKING: from collections.abc import AsyncIterator, Mapping @@ -15,7 +16,9 @@ from typing_extensions import Unpack - from crawlee._types import ExportDataCsvKwargs, ExportDataJsonKwargs, JsonSerializable + from crawlee._types import ExportDataCsvKwargs, ExportDataCsvWriterKwargs, ExportDataJsonKwargs, JsonSerializable + +logger = getLogger(__name__) if sys.platform == 'win32': @@ -196,23 +199,59 @@ async def export_csv_to_stream( dst: TextIO, **kwargs: Unpack[ExportDataCsvKwargs], ) -> None: - # Set lineterminator to '\n' if not explicitly provided. This prevents double line endings on Windows. - # The csv.writer default is '\r\n', which when written to a file in text mode on Windows gets converted - # to '\r\r\n' due to newline translation. By using '\n', we let the platform handle the line ending - # conversion: '\n' stays as '\n' on Unix, and becomes '\r\n' on Windows. + collect_all_keys = kwargs.pop('collect_all_keys', False) + + # Set lineterminator to '\n' if not explicitly provided. This prevents double line endings on Windows. The + # `csv.DictWriter` default is '\r\n', which when written to a file in text mode on Windows gets converted to + # '\r\r\n' due to newline translation. By using '\n', we let the platform handle the line ending conversion: + # '\n' stays as '\n' on Unix, and becomes '\r\n' on Windows. if 'lineterminator' not in kwargs: kwargs['lineterminator'] = '\n' - writer = csv.writer(dst, **kwargs) - write_header = True + writer_kwargs = cast('ExportDataCsvWriterKwargs', kwargs) + + # The real writer cannot be built until the first item's keys are known, so construct a throwaway one up front + # to validate the options. Without this, an export configured with e.g. `delimiter='ab'` would raise only when + # the dataset happens to contain an item, and silently succeed on a run that scraped nothing. + csv.DictWriter(dst, fieldnames=[], extrasaction='ignore', **writer_kwargs) + + if collect_all_keys: + items = [item async for item in iterator if item] + if not items: + return + + fieldnames = list(dict.fromkeys(key for item in items for key in item)) + writer = csv.DictWriter(dst, fieldnames=fieldnames, extrasaction='ignore', **writer_kwargs) + writer.writeheader() + for item in items: + writer.writerow(item) + return + + writer = None + header_keys: set[str] = set() + dropped_keys: set[str] = set() - # Iterate over the dataset and write to CSV. async for item in iterator: if not item: continue - if write_header: - writer.writerow(item.keys()) - write_header = False - - writer.writerow(item.values()) + if writer is None: + writer = csv.DictWriter(dst, fieldnames=list(item), extrasaction='ignore', **writer_kwargs) + writer.writeheader() + header_keys = set(writer.fieldnames) + + dropped_keys.update(item.keys() - header_keys) + writer.writerow(item) + + # The header comes from the first item, so any key introduced by a later item is dropped. Dropping them is + # intentional and matches Crawlee for JS. Doing it silently is not, so report it once for the whole export. + if dropped_keys: + logger.warning( + 'CSV export dropped %d key(s) not present in the first item: %s. ' + 'Pass collect_all_keys=True to include keys from all items as columns.', + len(dropped_keys), + # The keys are annotated as `str`, but nothing enforces that at runtime, and a storage client that + # does not round-trip items through JSON hands them over as pushed. Stringify them so building the + # message cannot fail on a key that is not a string, or on a set mixing several key types. + ', '.join(sorted(str(key) for key in dropped_keys)), + ) diff --git a/tests/unit/_utils/test_file.py b/tests/unit/_utils/test_file.py index f4989337e9..322ad3fe03 100644 --- a/tests/unit/_utils/test_file.py +++ b/tests/unit/_utils/test_file.py @@ -1,15 +1,19 @@ from __future__ import annotations from datetime import datetime, timezone -from typing import TYPE_CHECKING +from io import StringIO +from typing import TYPE_CHECKING, cast import pytest -from crawlee._utils.file import json_dumps, validate_subdirectory +from crawlee._utils.file import export_csv_to_stream, json_dumps, validate_subdirectory if TYPE_CHECKING: + from collections.abc import AsyncIterator, Mapping from pathlib import Path + from crawlee._types import JsonSerializable + async def test_json_dumps() -> None: assert await json_dumps({'key': 'value'}) == '{\n "key": "value"\n}' @@ -19,6 +23,210 @@ async def test_json_dumps() -> None: assert await json_dumps(datetime(2022, 1, 1, tzinfo=timezone.utc)) == '"2022-01-01 00:00:00+00:00"' +async def async_iter( + items: list[Mapping[str, JsonSerializable]], +) -> AsyncIterator[Mapping[str, JsonSerializable]]: + for item in items: + yield item + + +async def test_export_csv_to_stream_keeps_columns_aligned_for_heterogeneous_items() -> None: + """Values must be written under their own header column even when items have different key orders/sets.""" + dst = StringIO() + await export_csv_to_stream( + async_iter( + [ + {'name': 'Alice', 'age': 30}, + {'name': 'Bob', 'city': 'NYC', 'age': 25}, + {'age': 40, 'name': 'Carol'}, + ] + ), + dst, + lineterminator='\n', + ) + + assert dst.getvalue() == 'name,age\nAlice,30\nBob,25\nCarol,40\n' + + +async def test_export_csv_to_stream_collects_all_keys_when_requested() -> None: + """All item keys are included when key collection is enabled.""" + dst = StringIO() + await export_csv_to_stream( + async_iter([{'name': 'Alice', 'age': 30}, {'name': 'Bob', 'city': 'NYC', 'age': 25}]), + dst, + collect_all_keys=True, + lineterminator='\n', + ) + + assert dst.getvalue() == 'name,age,city\nAlice,30,\nBob,25,NYC\n' + + +@pytest.mark.parametrize( + 'collect_all_keys', + [ + pytest.param(False, id='header from first item'), + pytest.param(True, id='all keys collected'), + ], +) +async def test_export_csv_to_stream_skips_empty_items(*, collect_all_keys: bool) -> None: + """Empty mappings are skipped and neither define the header nor emit a blank row, in either column mode.""" + dst = StringIO() + await export_csv_to_stream( + async_iter([{}, {'id': 1, 'name': 'Item 1'}, {}, {'id': 2, 'name': 'Item 2'}]), + dst, + collect_all_keys=collect_all_keys, + lineterminator='\n', + ) + + assert dst.getvalue() == 'id,name\n1,Item 1\n2,Item 2\n' + + +async def test_export_csv_to_stream_handles_empty_iterator() -> None: + """An empty iterator produces no CSV content.""" + dst = StringIO() + await export_csv_to_stream(async_iter([]), dst) + + assert dst.getvalue() == '' + + +@pytest.mark.parametrize( + 'items', + [ + pytest.param([], id='no items'), + pytest.param([{}, {}], id='only empty items'), + ], +) +async def test_export_csv_to_stream_collects_all_keys_without_writable_items( + items: list[Mapping[str, JsonSerializable]], +) -> None: + """Key collection produces no CSV content, not a bare header line, when there is nothing to write.""" + dst = StringIO() + await export_csv_to_stream(async_iter(items), dst, collect_all_keys=True) + + assert dst.getvalue() == '' + + +async def test_export_csv_to_stream_writes_before_consuming_all_items() -> None: + """The default export writes its header before requesting the second item.""" + dst = StringIO() + + async def items() -> AsyncIterator[Mapping[str, JsonSerializable]]: + yield {'id': 1} + assert dst.getvalue() == 'id\n1\n' + yield {'id': 2} + + await export_csv_to_stream(items(), dst, lineterminator='\n') + + assert dst.getvalue() == 'id\n1\n2\n' + + +async def test_export_csv_to_stream_preserves_non_json_values() -> None: + """CSV values do not pass through JSON serialization.""" + dst = StringIO() + value = datetime(2020, 1, 1, tzinfo=timezone.utc) + + await export_csv_to_stream( + async_iter([{'created_at': cast('JsonSerializable', value)}]), + dst, + lineterminator='\n', + ) + + assert dst.getvalue() == 'created_at\n2020-01-01 00:00:00+00:00\n' + + +async def test_export_csv_to_stream_warns_about_dropped_keys(caplog: pytest.LogCaptureFixture) -> None: + """Keys dropped because they are absent from the first item are reported once, naming every dropped key.""" + dst = StringIO() + with caplog.at_level('WARNING', logger='crawlee._utils.file'): + await export_csv_to_stream( + async_iter([{'name': 'Alice'}, {'name': 'Bob', 'city': 'NYC'}, {'name': 'Carol', 'age': 40}]), + dst, + lineterminator='\n', + ) + + assert dst.getvalue() == 'name\nAlice\nBob\nCarol\n' + assert len(caplog.records) == 1 + assert 'age, city' in caplog.records[0].message + assert 'collect_all_keys=True' in caplog.records[0].message + + +async def test_export_csv_to_stream_warns_about_dropped_keys_that_are_not_strings( + caplog: pytest.LogCaptureFixture, +) -> None: + """Dropped keys are reported even when they are not strings, which a storage client may hand over as pushed.""" + dst = StringIO() + items = cast('list[Mapping[str, JsonSerializable]]', [{'name': 'Alice'}, {1: 'x', 'city': 'NYC'}]) + + with caplog.at_level('WARNING', logger='crawlee._utils.file'): + await export_csv_to_stream(async_iter(items), dst, lineterminator='\n') + + # The second item has no value for the only column, so its row holds a single empty field, written as '""' to + # keep it distinguishable from a blank line. + assert dst.getvalue() == 'name\nAlice\n""\n' + assert len(caplog.records) == 1 + assert '1, city' in caplog.records[0].message + + +async def test_export_csv_to_stream_does_not_warn_when_nothing_is_dropped( + caplog: pytest.LogCaptureFixture, +) -> None: + """Collecting all keys drops nothing, so it must not warn even for items with differing key sets.""" + dst = StringIO() + with caplog.at_level('WARNING', logger='crawlee._utils.file'): + await export_csv_to_stream( + async_iter([{'name': 'Alice'}, {'name': 'Bob', 'city': 'NYC'}]), + dst, + collect_all_keys=True, + lineterminator='\n', + ) + + assert dst.getvalue() == 'name,city\nAlice,\nBob,NYC\n' + assert caplog.records == [] + + +@pytest.mark.parametrize( + 'collect_all_keys', + [ + pytest.param(False, id='header from first item'), + pytest.param(True, id='all keys collected'), + ], +) +async def test_export_csv_to_stream_rejects_invalid_writer_options_for_empty_iterator( + *, collect_all_keys: bool +) -> None: + """Writer options are validated up front, so a misconfigured export fails even when there is nothing to write.""" + # Match only the option name: the rest of the message is CPython's wording and it changed in 3.14. + with pytest.raises(TypeError, match='"delimiter"'): + await export_csv_to_stream(async_iter([]), StringIO(), delimiter='ab', collect_all_keys=collect_all_keys) + + +async def test_export_csv_to_stream_honors_restval() -> None: + """`restval` fills the cells of columns an item does not contain.""" + dst = StringIO() + await export_csv_to_stream( + async_iter([{'name': 'Alice'}, {'name': 'Bob', 'city': 'NYC'}]), + dst, + collect_all_keys=True, + restval='N/A', + lineterminator='\n', + ) + + assert dst.getvalue() == 'name,city\nAlice,N/A\nBob,NYC\n' + + +async def test_export_csv_to_stream_honors_restval_without_collecting_all_keys() -> None: + """`restval` also fills the cells of header columns a later item does not contain in the default column mode.""" + dst = StringIO() + await export_csv_to_stream( + async_iter([{'name': 'Alice', 'city': 'NYC'}, {'name': 'Bob'}]), + dst, + restval='N/A', + lineterminator='\n', + ) + + assert dst.getvalue() == 'name,city\nAlice,NYC\nBob,N/A\n' + + # Tests for validate_subdirectory (storage name/alias directory validation). diff --git a/tests/unit/crawlers/_basic/test_basic_crawler.py b/tests/unit/crawlers/_basic/test_basic_crawler.py index b8c33d2eaf..763969b3f9 100644 --- a/tests/unit/crawlers/_basic/test_basic_crawler.py +++ b/tests/unit/crawlers/_basic/test_basic_crawler.py @@ -783,6 +783,22 @@ async def test_crawler_export_data_additional_kwargs(tmp_path: Path) -> None: assert csv_path.read_text() == 'z;a\n1;2\n' +async def test_crawler_export_data_csv_collect_all_keys(tmp_path: Path) -> None: + crawler = BasicCrawler() + dataset = await Dataset.open() + + await dataset.push_data([{'a': 1}, {'a': 2, 'b': 3}]) + + default_path = tmp_path / 'default.csv' + all_keys_path = tmp_path / 'all_keys.csv' + + await crawler.export_data(path=default_path, lineterminator='\n') + await crawler.export_data(path=all_keys_path, collect_all_keys=True, lineterminator='\n') + + assert default_path.read_text() == 'a\n1\n2\n' + assert all_keys_path.read_text() == 'a,b\n1,\n2,3\n' + + async def test_context_push_and_export_data(tmp_path: Path) -> None: crawler = BasicCrawler() diff --git a/tests/unit/storages/test_dataset.py b/tests/unit/storages/test_dataset.py index 6c630552f3..c18c3e54f3 100644 --- a/tests/unit/storages/test_dataset.py +++ b/tests/unit/storages/test_dataset.py @@ -472,6 +472,46 @@ async def test_export_to_csv( await kvs.drop() +async def test_export_to_csv_columns( + dataset: Dataset, + storage_client: StorageClient, +) -> None: + """Test the CSV columns come from the first item by default, and from all items with `collect_all_keys`.""" + kvs = await KeyValueStore.open( + name='export-kvs', + storage_client=storage_client, + ) + + await dataset.push_data( + [ + {'name': 'Alice', 'age': 30}, + {'name': 'Bob', 'city': 'NYC', 'age': 25}, + {'age': 40, 'name': 'Carol'}, + ] + ) + + await dataset.export_to( + key='first_item_keys.csv', + content_type='csv', + to_kvs_name='export-kvs', + to_kvs_storage_client=storage_client, + lineterminator='\n', + ) + await dataset.export_to( + key='all_keys.csv', + content_type='csv', + to_kvs_name='export-kvs', + to_kvs_storage_client=storage_client, + collect_all_keys=True, + lineterminator='\n', + ) + + assert await kvs.get_value(key='first_item_keys.csv') == 'name,age\nAlice,30\nBob,25\nCarol,40\n' + assert await kvs.get_value(key='all_keys.csv') == 'name,age,city\nAlice,30,\nBob,25,NYC\nCarol,40,\n' + + await kvs.drop() + + async def test_export_to_invalid_content_type(dataset: Dataset) -> None: """Test exporting dataset with invalid content type raises error.""" with pytest.raises(ValueError, match=r'Unsupported content type'):