Address feedback

Author: vdusek

src/crawlee/_autoscaling/autoscaled_pool.py

@@ -142,8 +142,7 @@ async def run(self) -> None:
 
             logger.info('Waiting for remaining tasks to finish')
 
-            tasks_to_wait = list(run.worker_tasks)
-            for task in tasks_to_wait:
+            for task in run.worker_tasks:
                 if not task.done():
                     with suppress(BaseException):
                         await task

src/crawlee/_utils/file.py

@@ -6,7 +6,7 @@
 import os
 import tempfile
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, overload
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
@@ -54,59 +54,75 @@ async def json_dumps(obj: Any) -> str:
     return await asyncio.to_thread(json.dumps, obj, ensure_ascii=False, indent=2, default=str)
 
 
-async def atomic_write_text(path: Path, data: str) -> None:
-    dir_path = path.parent
+@overload
+async def atomic_write(
+    path: Path,
+    data: str,
+    *,
+    is_binary: bool = False,
+) -> None: ...
 
-    def _sync_write_text() -> str:
-        # create a temp file in the target dir, return its name
-        fd, tmp_path = tempfile.mkstemp(
-            suffix=f'{path.suffix}.tmp',
-            prefix=f'{path.name}.',
-            dir=str(dir_path),
-        )
-        try:
-            with os.fdopen(fd, 'w', encoding='utf-8') as tmp_file:
-                tmp_file.write(data)
-        except:
-            Path(tmp_path).unlink(missing_ok=True)
-            raise
-        return tmp_path
 
-    tmp_path = await asyncio.to_thread(_sync_write_text)
+@overload
+async def atomic_write(
+    path: Path,
+    data: bytes,
+    *,
+    is_binary: bool = True,
+) -> None: ...
 
-    try:
-        await asyncio.to_thread(os.replace, tmp_path, str(path))
-    except (FileNotFoundError, PermissionError):
-        # fallback if tmp went missing
-        await asyncio.to_thread(path.write_text, data, encoding='utf-8')
-    finally:
-        await asyncio.to_thread(Path(tmp_path).unlink, missing_ok=True)
 
+async def atomic_write(
+    path: Path,
+    data: str | bytes,
+    *,
+    is_binary: bool = False,
+) -> None:
+    """Write data to a file atomically to prevent data corruption or partial writes.
 
-async def atomic_write_bytes(path: Path, data: bytes) -> None:
+    This function handles both text and binary data. It ensures atomic writing by creating
+    a temporary file and then atomically replacing the target file, which prevents data
+    corruption if the process is interrupted during the write operation.
+
+    For example, if a process (crawler) is interrupted while writing a file, the file may end up in an
+    incomplete or corrupted state. This might be especially unwanted for metadata files.
+
+    Args:
+        path: The path to the destination file.
+        data: The data to write to the file (string or bytes).
+        is_binary: If True, write in binary mode. If False (default), write in text mode.
+    """
     dir_path = path.parent
 
-    def _sync_write_bytes() -> str:
+    def _sync_write() -> str:
+        # create a temp file in the target dir, return its name
         fd, tmp_path = tempfile.mkstemp(
             suffix=f'{path.suffix}.tmp',
             prefix=f'{path.name}.',
             dir=str(dir_path),
         )
         try:
-            with os.fdopen(fd, 'wb') as tmp_file:
-                tmp_file.write(data)
-        except:
+            if is_binary:
+                with os.fdopen(fd, 'wb') as tmp_file:
+                    tmp_file.write(data)  # type: ignore[arg-type]
+            else:
+                with os.fdopen(fd, 'w', encoding='utf-8') as tmp_file:
+                    tmp_file.write(data)  # type: ignore[arg-type]
+        except Exception:  # broader exception handling
             Path(tmp_path).unlink(missing_ok=True)
             raise
         return tmp_path
 
-    tmp_path = await asyncio.to_thread(_sync_write_bytes)
+    tmp_path = await asyncio.to_thread(_sync_write)
 
     try:
         await asyncio.to_thread(os.replace, tmp_path, str(path))
     except (FileNotFoundError, PermissionError):
         # fallback if tmp went missing
-        await asyncio.to_thread(path.write_bytes, data)
+        if is_binary:
+            await asyncio.to_thread(path.write_bytes, data)  # type: ignore[arg-type]
+        else:
+            await asyncio.to_thread(path.write_text, data, encoding='utf-8')  # type: ignore[arg-type]
     finally:
         await asyncio.to_thread(Path(tmp_path).unlink, missing_ok=True)
 

src/crawlee/statistics/_error_snapshotter.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import hashlib
 import re
 import string
@@ -28,23 +29,42 @@ async def capture_snapshot(
         file_and_line: str,
         context: BasicCrawlingContext,
     ) -> None:
-        """Capture error snapshot and save it to key value store."""
-        snapshot = await context.get_snapshot()
-        if not snapshot:
-            return
-
-        base = self._get_snapshot_base_name(error_message, file_and_line)
-        kvs = await KeyValueStore.open(name=self._kvs_name)
-
-        # Save HTML snapshot if present
-        if snapshot.html:
-            key_html = f'{base}.html'
-            await kvs.set_value(key_html, snapshot.html, content_type='text/html')
-
-        # Save screenshot snapshot if present
-        if snapshot.screenshot:
-            key_jpg = f'{base}.jpg'
-            await kvs.set_value(key_jpg, snapshot.screenshot, content_type='image/jpeg')
+        """Capture error snapshot and save it to key value store.
+
+        It saves the error snapshot directly to a key value store. It can't use `context.get_key_value_store` because
+        it returns `KeyValueStoreChangeRecords` which is commited to the key value store only if the `RequestHandler`
+        returned without an exception. ErrorSnapshotter is on the contrary active only when `RequestHandler` fails with
+        an exception.
+
+        Args:
+            error_message: Used in filename of the snapshot.
+            file_and_line: Used in filename of the snapshot.
+            context: Context that is used to get the snapshot.
+        """
+        if snapshot := await context.get_snapshot():
+            kvs = await KeyValueStore.open(name=self._kvs_name)
+            snapshot_base_name = self._get_snapshot_base_name(error_message, file_and_line)
+            snapshot_save_tasks = list[asyncio.Task]()
+
+            if snapshot.html:
+                snapshot_save_tasks.append(
+                    asyncio.create_task(self._save_html(kvs, snapshot.html, base_name=snapshot_base_name))
+                )
+
+            if snapshot.screenshot:
+                snapshot_save_tasks.append(
+                    asyncio.create_task(self._save_screenshot(kvs, snapshot.screenshot, base_name=snapshot_base_name))
+                )
+
+            await asyncio.gather(*snapshot_save_tasks)
+
+    async def _save_html(self, kvs: KeyValueStore, html: str, base_name: str) -> None:
+        file_name = f'{base_name}.html'
+        await kvs.set_value(file_name, html, content_type='text/html')
+
+    async def _save_screenshot(self, kvs: KeyValueStore, screenshot: bytes, base_name: str) -> None:
+        file_name = f'{base_name}.jpg'
+        await kvs.set_value(file_name, screenshot, content_type='image/jpeg')
 
     def _sanitize_filename(self, filename: str) -> str:
         return re.sub(f'[^{re.escape(self.ALLOWED_CHARACTERS)}]', '', filename[: self.MAX_FILENAME_LENGTH])

src/crawlee/storage_clients/_file_system/_dataset_client.py

@@ -13,7 +13,7 @@
 
 from crawlee._consts import METADATA_FILENAME
 from crawlee._utils.crypto import crypto_random_object_id
-from crawlee._utils.file import atomic_write_text, json_dumps
+from crawlee._utils.file import atomic_write, json_dumps
 from crawlee.storage_clients._base import DatasetClient
 from crawlee.storage_clients.models import DatasetItemsListPage, DatasetMetadata
 
@@ -442,7 +442,7 @@ async def _update_metadata(
 
         # Dump the serialized metadata to the file.
         data = await json_dumps(self._metadata.model_dump())
-        await atomic_write_text(self.path_to_metadata, data)
+        await atomic_write(self.path_to_metadata, data)
 
     async def _push_item(self, item: dict[str, Any], item_id: int) -> None:
         """Push a single item to the dataset.
@@ -463,7 +463,7 @@ async def _push_item(self, item: dict[str, Any], item_id: int) -> None:
 
         # Dump the serialized item to the file.
         data = await json_dumps(item)
-        await atomic_write_text(file_path, data)
+        await atomic_write(file_path, data)
 
     async def _get_sorted_data_files(self) -> list[Path]:
         """Retrieve and return a sorted list of data files in the dataset directory.

src/crawlee/storage_clients/_file_system/_key_value_store_client.py

@@ -14,7 +14,7 @@
 
 from crawlee._consts import METADATA_FILENAME
 from crawlee._utils.crypto import crypto_random_object_id
-from crawlee._utils.file import atomic_write_bytes, atomic_write_text, infer_mime_type, json_dumps
+from crawlee._utils.file import atomic_write, infer_mime_type, json_dumps
 from crawlee.storage_clients._base import KeyValueStoreClient
 from crawlee.storage_clients.models import KeyValueStoreMetadata, KeyValueStoreRecord, KeyValueStoreRecordMetadata
 
@@ -325,10 +325,10 @@ async def set_value(self, *, key: str, value: Any, content_type: str | None = No
             await asyncio.to_thread(self.path_to_kvs.mkdir, parents=True, exist_ok=True)
 
             # Write the value to the file.
-            await atomic_write_bytes(record_path, value_bytes)
+            await atomic_write(record_path, value_bytes, is_binary=True)
 
             # Write the record metadata to the file.
-            await atomic_write_text(record_metadata_filepath, record_metadata_content)
+            await atomic_write(record_metadata_filepath, record_metadata_content)
 
             # Update the KVS metadata to record the access and modification.
             await self._update_metadata(update_accessed_at=True, update_modified_at=True)
@@ -445,7 +445,7 @@ async def _update_metadata(
 
         # Dump the serialized metadata to the file.
         data = await json_dumps(self._metadata.model_dump())
-        await atomic_write_text(self.path_to_metadata, data)
+        await atomic_write(self.path_to_metadata, data)
 
     def _encode_key(self, key: str) -> str:
         """Encode a key to make it safe for use in a file path."""

src/crawlee/storage_clients/_file_system/_request_queue_client.py

@@ -14,7 +14,7 @@
 from crawlee import Request
 from crawlee._consts import METADATA_FILENAME
 from crawlee._utils.crypto import crypto_random_object_id
-from crawlee._utils.file import atomic_write_text, json_dumps
+from crawlee._utils.file import atomic_write, json_dumps
 from crawlee.storage_clients._base import RequestQueueClient
 from crawlee.storage_clients.models import AddRequestsResponse, ProcessedRequest, RequestQueueMetadata
 
@@ -336,7 +336,7 @@ async def add_batch_of_requests(
                     # Update the existing request file
                     request_path = self.path_to_rq / f'{existing_request.id}.json'
                     request_data = await json_dumps(existing_request.model_dump())
-                    await atomic_write_text(request_path, request_data)
+                    await atomic_write(request_path, request_data)
 
                     processed_requests.append(
                         ProcessedRequest(
@@ -362,7 +362,7 @@ async def add_batch_of_requests(
                 request_dict['_sequence'] = sequence_number
 
                 request_data = await json_dumps(request_dict)
-                await atomic_write_text(request_path, request_data)
+                await atomic_write(request_path, request_data)
 
                 # Update metadata counts
                 new_total_request_count += 1
@@ -620,7 +620,7 @@ async def mark_request_as_handled(self, request: Request) -> ProcessedRequest |
                 return None
 
             request_data = await json_dumps(request.model_dump())
-            await atomic_write_text(request_path, request_data)
+            await atomic_write(request_path, request_data)
 
             # Update metadata timestamps
             await self._update_metadata(
@@ -678,7 +678,7 @@ async def reclaim_request(
                 return None
 
             request_data = await json_dumps(request.model_dump())
-            await atomic_write_text(request_path, request_data)
+            await atomic_write(request_path, request_data)
 
             # Update metadata timestamps
             await self._update_metadata(update_modified_at=True, update_accessed_at=True)
@@ -781,4 +781,4 @@ async def _update_metadata(
 
         # Dump the serialized metadata to the file.
         data = await json_dumps(self._metadata.model_dump())
-        await atomic_write_text(self.path_to_metadata, data)
+        await atomic_write(self.path_to_metadata, data)

tests/unit/crawlers/_basic/test_basic_crawler.py

@@ -618,7 +618,6 @@ async def test_crawler_get_storages() -> None:
     assert isinstance(kvs, KeyValueStore)
 
 
-# THIS
 async def test_crawler_run_requests() -> None:
     crawler = BasicCrawler()
     seen_urls = list[str]()

tests/unit/storage_clients/_file_system/test_fs_kvs_client.py

@@ -312,9 +312,9 @@ async def test_iterate_keys_with_exclusive_start_key(kvs_client: FileSystemKeyVa
     """Test that exclusive_start_key parameter returns only keys after it alphabetically."""
     # Add some values with alphabetical keys
     await kvs_client.set_value(key='a-key', value='value-a')
-    await kvs_client.set_value(key='b-key', value='value-b')
-    await kvs_client.set_value(key='c-key', value='value-c')
     await kvs_client.set_value(key='d-key', value='value-d')
+    await kvs_client.set_value(key='c-key', value='value-c')
+    await kvs_client.set_value(key='b-key', value='value-b')
 
     # Iterate with exclusive start key
     keys = [key.key async for key in kvs_client.iterate_keys(exclusive_start_key='b-key')]

tests/unit/storage_clients/_memory/test_memory_kvs_client.py

@@ -167,7 +167,7 @@ async def test_iterate_keys(kvs_client: MemoryKeyValueStoreClient) -> None:
 async def test_iterate_keys_with_exclusive_start_key(kvs_client: MemoryKeyValueStoreClient) -> None:
     """Test that exclusive_start_key parameter returns only keys after it alphabetically."""
     # Set some values
-    for key in ['a_key', 'b_key', 'c_key', 'd_key', 'e_key']:
+    for key in ['b_key', 'c_key', 'a_key', 'e_key', 'd_key']:
         await kvs_client.set_value(key=key, value=f'value for {key}')
 
     # Get keys starting after 'b_key'
@@ -181,7 +181,7 @@ async def test_iterate_keys_with_exclusive_start_key(kvs_client: MemoryKeyValueS
 async def test_iterate_keys_with_limit(kvs_client: MemoryKeyValueStoreClient) -> None:
     """Test that the limit parameter returns only the specified number of keys."""
     # Set some values
-    for key in ['a_key', 'b_key', 'c_key', 'd_key', 'e_key']:
+    for key in ['a_key', 'e_key', 'c_key', 'b_key', 'd_key']:
         await kvs_client.set_value(key=key, value=f'value for {key}')
 
     # Get first 3 keys