feat(dataset): add collect_all_keys option and keep CSV export columns aligned

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 <ApiLink to="class/BasicCrawler#export_data">`BasicCrawler.export_data`</ApiLink> 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 <ApiLink to="class/BasicCrawler#export_data">`BasicCrawler.export_data`</ApiLink> 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 <ApiLink to="class/BeautifulSoupCrawler">`B
         </RunnableCodeBlock>
     </TabItem>
 </Tabs>
+
+## 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.

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,13 +804,25 @@ 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."""
 
     strict: NotRequired[bool]
     """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`.
 

src/crawlee/_utils/file.py

@@ -6,16 +6,19 @@
 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
     from typing import Any, TextIO
 
     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)),
+        )

src/crawlee/crawlers/_basic/_basic_crawler.py

@@ -111,6 +111,13 @@
 FailedRequestHandler = Callable[[TCrawlingContext, Exception], Awaitable[None]]
 SkippedRequestCallback = Callable[[str, SkippedReason], Awaitable[None]]
 
+# `export_data` accepts the kwargs of both export formats, since the format is only known once the destination path
+# is inspected. These are used to forward just the ones the selected format understands. Both key sets are unioned,
+# because `NotRequired` is invisible to `TypedDict` under postponed annotation evaluation, which makes every key land
+# in `__required_keys__` instead of `__optional_keys__`.
+_CSV_EXPORT_KEYS = ExportDataCsvKwargs.__required_keys__ | ExportDataCsvKwargs.__optional_keys__
+_JSON_EXPORT_KEYS = ExportDataJsonKwargs.__required_keys__ | ExportDataJsonKwargs.__optional_keys__
+
 
 class _BasicCrawlerOptions(TypedDict):
     """Non-generic options the `BasicCrawler` constructor."""
@@ -940,12 +947,18 @@ async def export_data(
 
         if path.suffix == '.csv':
             dst = StringIO()
-            csv_kwargs = cast('ExportDataCsvKwargs', additional_kwargs)
+            csv_kwargs = cast(
+                'ExportDataCsvKwargs',
+                {key: value for key, value in additional_kwargs.items() if key in _CSV_EXPORT_KEYS},
+            )
             await export_csv_to_stream(dataset.iterate_items(), dst, **csv_kwargs)
             await atomic_write(path, dst.getvalue())
         elif path.suffix == '.json':
             dst = StringIO()
-            json_kwargs = cast('ExportDataJsonKwargs', additional_kwargs)
+            json_kwargs = cast(
+                'ExportDataJsonKwargs',
+                {key: value for key, value in additional_kwargs.items() if key in _JSON_EXPORT_KEYS},
+            )
             await export_json_to_stream(dataset.iterate_items(), dst, **json_kwargs)
             await atomic_write(path, dst.getvalue())
         else: