feat: BasicCrawler.export_data helper (#222)

- closes #211

Author: janbuchar

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Adds
 
+- BREAKING: `BasicCrawler.export_data` helper method which replaces `BasicCrawler.export_to`
 - `Configuration.get_global_configuration` method
 
 ### Fixes

src/crawlee/basic_crawler/basic_crawler.py

@@ -11,7 +11,8 @@
 from datetime import timedelta
 from functools import partial
 from logging import getLogger
-from typing import TYPE_CHECKING, Any, AsyncContextManager, Callable, Generic, Union, cast
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, AsyncContextManager, Callable, Generic, Literal, Union, cast
 
 import httpx
 from tldextract import TLDExtract
@@ -48,7 +49,7 @@
     from crawlee.proxy_configuration import ProxyConfiguration, ProxyInfo
     from crawlee.sessions import Session
     from crawlee.statistics import FinalStatistics, StatisticsState
-    from crawlee.storages.dataset import ExportToKwargs, GetDataKwargs, PushDataKwargs
+    from crawlee.storages.dataset import GetDataKwargs, PushDataKwargs
     from crawlee.storages.request_provider import RequestProvider
     from crawlee.types import JSONSerializable
 
@@ -411,24 +412,31 @@ async def get_data(
         dataset = await Dataset.open(id=dataset_id, name=dataset_name)
         return await dataset.get_data(**kwargs)
 
-    async def export_to(
+    async def export_data(
         self,
+        path: str | Path,
+        content_type: Literal['json', 'csv'] | None = None,
         dataset_id: str | None = None,
         dataset_name: str | None = None,
-        **kwargs: Unpack[ExportToKwargs],
     ) -> None:
         """Export data from a dataset.
 
         This helper method simplifies the process of exporting data from a dataset. It opens the specified
         dataset and then exports the data based on the provided parameters.
 
         Args:
+            path: The destination path
+            content_type: The output format
             dataset_id: The ID of the dataset.
             dataset_name: The name of the dataset.
-            kwargs: Keyword arguments to be passed to the dataset's `export_to` method.
         """
         dataset = await Dataset.open(id=dataset_id, name=dataset_name)
-        return await dataset.export_to(**kwargs)
+        path = path if isinstance(path, Path) else Path(path)
+
+        if content_type is None:
+            content_type = 'csv' if path.suffix == '.csv' else 'json'
+
+        return await dataset.write_to(content_type, path.open('w', newline=''))
 
     async def _push_data(
         self,

src/crawlee/storages/dataset.py

@@ -2,7 +2,8 @@
 
 import csv
 import io
-from typing import TYPE_CHECKING, AsyncIterator, Literal, TypedDict, cast
+import json
+from typing import TYPE_CHECKING, AsyncIterator, Literal, TextIO, TypedDict, cast
 
 from typing_extensions import NotRequired, Required, Unpack, override
 
@@ -189,6 +190,32 @@ async def get_data(self, **kwargs: Unpack[GetDataKwargs]) -> DatasetItemsListPag
         # https://github.com/apify/apify-sdk-python/issues/140
         return await self._resource_client.list_items(**kwargs)
 
+    async def write_to(self, content_type: Literal['json', 'csv'], destination: TextIO) -> None:
+        """Exports the entire dataset into an arbitrary stream.
+
+        Args:
+            content_type: Specifies the output format
+            destination: The stream into which the dataset contents should be written
+        """
+        items: list[dict] = []
+        limit = 1000
+        offset = 0
+
+        while True:
+            list_items = await self._resource_client.list_items(limit=limit, offset=offset)
+            items.extend(list_items.items)
+            if list_items.total <= offset + list_items.count:
+                break
+            offset += list_items.count
+
+        if content_type == 'csv':
+            writer = csv.writer(destination, quoting=csv.QUOTE_MINIMAL)
+            writer.writerows([items[0].keys(), *[item.values() for item in items]])
+        elif content_type == 'json':
+            json.dump(items, destination)
+        else:
+            raise ValueError(f'Unsupported content type: {content_type}')
+
     async def export_to(self, **kwargs: Unpack[ExportToKwargs]) -> None:
         """Exports the entire dataset into a specified file stored under a key in a key-value store.
 
@@ -206,30 +233,15 @@ async def export_to(self, **kwargs: Unpack[ExportToKwargs]) -> None:
         to_key_value_store_name = kwargs.get('to_key_value_store_name', None)
 
         key_value_store = await KeyValueStore.open(id=to_key_value_store_id, name=to_key_value_store_name)
-        items: list[dict] = []
-        limit = 1000
-        offset = 0
 
-        while True:
-            list_items = await self._resource_client.list_items(limit=limit, offset=offset)
-            items.extend(list_items.items)
-            if list_items.total <= offset + list_items.count:
-                break
-            offset += list_items.count
+        output = io.StringIO()
+        await self.write_to(content_type, output)
 
         if content_type == 'csv':
-            content_type_full = 'text/csv'
-            output = io.StringIO()
-            writer = csv.writer(output, quoting=csv.QUOTE_MINIMAL)
-            writer.writerows([items[0].keys(), *[item.values() for item in items]])
-            value = output.getvalue()
-            return await key_value_store.set_value(key, value, content_type_full)
+            await key_value_store.set_value(key, output.getvalue(), 'text/csv')
 
         if content_type == 'json':
-            content_type_full = 'application/json'
-            return await key_value_store.set_value(key, items, content_type_full)
-
-        raise ValueError(f'Unsupported content type: {content_type}')
+            await key_value_store.set_value(key, output.getvalue(), 'application/json')
 
     async def get_info(self) -> DatasetMetadata | None:
         """Get an object containing general information about the dataset."""

tests/unit/basic_crawler/test_basic_crawler.py

@@ -23,6 +23,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
+    from pathlib import Path
 
     import respx
 
@@ -516,26 +517,25 @@ async def handler(context: BasicCrawlingContext) -> None:
     assert stats.requests_finished == 1
 
 
-async def test_crawler_push_and_export_data() -> None:
+async def test_crawler_push_and_export_data(tmp_path: Path) -> None:
     crawler = BasicCrawler()
     dataset = await Dataset.open()
 
     await dataset.push_data([{'id': 0, 'test': 'test'}, {'id': 1, 'test': 'test'}])
     await dataset.push_data({'id': 2, 'test': 'test'})
 
-    await crawler.export_to(key='dataset-json', content_type='json')
-    await crawler.export_to(key='dataset-csv', content_type='csv')
+    await crawler.export_data(tmp_path / 'dataset.json')
+    await crawler.export_data(tmp_path / 'dataset.csv')
 
-    kvs = await KeyValueStore.open()
-    assert await kvs.get_value('dataset-json') == [
+    assert json.load((tmp_path / 'dataset.json').open()) == [
         {'id': 0, 'test': 'test'},
         {'id': 1, 'test': 'test'},
         {'id': 2, 'test': 'test'},
     ]
-    assert await kvs.get_value('dataset-csv') == 'id,test\r\n0,test\r\n1,test\r\n2,test\r\n'
+    assert (tmp_path / 'dataset.csv').read_bytes() == b'id,test\r\n0,test\r\n1,test\r\n2,test\r\n'
 
 
-async def test_context_push_and_export_data(httpbin: str) -> None:
+async def test_context_push_and_export_data(httpbin: str, tmp_path: Path) -> None:
     crawler = BasicCrawler()
 
     @crawler.router.default_handler
@@ -545,16 +545,16 @@ async def handler(context: BasicCrawlingContext) -> None:
 
     await crawler.run([f'{httpbin}/1'])
 
-    await crawler.export_to(key='dataset-json', content_type='json')
-    await crawler.export_to(key='dataset-csv', content_type='csv')
+    await crawler.export_data(tmp_path / 'dataset.json')
+    await crawler.export_data(tmp_path / 'dataset.csv')
 
-    kvs = await KeyValueStore.open()
-    assert await kvs.get_value('dataset-json') == [
+    assert json.load((tmp_path / 'dataset.json').open()) == [
         {'id': 0, 'test': 'test'},
         {'id': 1, 'test': 'test'},
         {'id': 2, 'test': 'test'},
     ]
-    assert await kvs.get_value('dataset-csv') == 'id,test\r\n0,test\r\n1,test\r\n2,test\r\n'
+
+    assert (tmp_path / 'dataset.csv').read_bytes() == b'id,test\r\n0,test\r\n1,test\r\n2,test\r\n'
 
 
 async def test_max_requests_per_crawl(httpbin: str) -> None: