Update storage clients

Author: vdusek

src/crawlee/_service_locator.py

@@ -3,8 +3,8 @@
 from crawlee._utils.docs import docs_group
 from crawlee.configuration import Configuration
 from crawlee.errors import ServiceConflictError
-from crawlee.events import EventManager
-from crawlee.storage_clients import StorageClient
+from crawlee.events import EventManager, LocalEventManager
+from crawlee.storage_clients import FileSystemStorageClient, StorageClient
 
 
 @docs_group('Classes')
@@ -49,8 +49,6 @@ def set_configuration(self, configuration: Configuration) -> None:
     def get_event_manager(self) -> EventManager:
         """Get the event manager."""
         if self._event_manager is None:
-            from crawlee.events import LocalEventManager
-
             self._event_manager = (
                 LocalEventManager().from_config(config=self._configuration)
                 if self._configuration
@@ -77,9 +75,7 @@ def set_event_manager(self, event_manager: EventManager) -> None:
     def get_storage_client(self) -> StorageClient:
         """Get the storage client."""
         if self._storage_client is None:
-            from crawlee.storage_clients import file_system_storage_client
-
-            self._storage_client = file_system_storage_client
+            self._storage_client = FileSystemStorageClient()
 
         self._storage_client_was_retrieved = True
         return self._storage_client

src/crawlee/storage_clients/__init__.py

@@ -1,9 +1,9 @@
 from ._base import StorageClient
-from ._file_system import file_system_storage_client
-from ._memory import memory_storage_client
+from ._file_system import FileSystemStorageClient
+from ._memory import MemoryStorageClient
 
 __all__ = [
+    'FileSystemStorageClient',
+    'MemoryStorageClient',
     'StorageClient',
-    'file_system_storage_client',
-    'memory_storage_client'
 ]

src/crawlee/storage_clients/_base/_storage_client.py

@@ -1,16 +1,48 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from ._dataset_client import DatasetClient
     from ._key_value_store_client import KeyValueStoreClient
     from ._request_queue_client import RequestQueueClient
 
 
-@dataclass
-class StorageClient:
-    dataset_client_class: type[DatasetClient]
-    key_value_store_client_class: type[KeyValueStoreClient]
-    request_queue_client_class: type[RequestQueueClient]
+class StorageClient(ABC):
+    """Base class for storage clients."""
+
+    @abstractmethod
+    async def open_dataset_client(
+        self,
+        *,
+        id: str | None,
+        name: str | None,
+        purge_on_start: bool,
+        storage_dir: Path,
+    ) -> DatasetClient:
+        """Open the dataset client."""
+
+    @abstractmethod
+    async def open_key_value_store_client(
+        self,
+        *,
+        id: str | None,
+        name: str | None,
+        purge_on_start: bool,
+        storage_dir: Path,
+    ) -> KeyValueStoreClient:
+        """Open the key-value store client."""
+
+    @abstractmethod
+    async def open_request_queue_client(
+        self,
+        *,
+        id: str | None,
+        name: str | None,
+        purge_on_start: bool,
+        storage_dir: Path,
+    ) -> RequestQueueClient:
+        """Open the request queue client."""

src/crawlee/storage_clients/_file_system/__init__.py

@@ -1,3 +1,3 @@
-from ._storage_client import file_system_storage_client
+from ._storage_client import FileSystemStorageClient
 
-__all__ = ['file_system_storage_client']
+__all__ = ['FileSystemStorageClient']

src/crawlee/storage_clients/_file_system/_storage_client.py

@@ -1,13 +1,65 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
+from typing_extensions import override
+
 from crawlee.storage_clients._base import StorageClient
 
 from ._dataset_client import FileSystemDatasetClient
 from ._key_value_store import FileSystemKeyValueStoreClient
 from ._request_queue import FileSystemRequestQueueClient
 
-file_system_storage_client = StorageClient(
-    dataset_client_class=FileSystemDatasetClient,
-    key_value_store_client_class=FileSystemKeyValueStoreClient,
-    request_queue_client_class=FileSystemRequestQueueClient,
-)
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+class FileSystemStorageClient(StorageClient):
+    """File system storage client."""
+
+    @override
+    async def open_dataset_client(
+        self,
+        *,
+        id: str | None,
+        name: str | None,
+        purge_on_start: bool,
+        storage_dir: Path,
+    ) -> FileSystemDatasetClient:
+        dataset_client = await FileSystemDatasetClient.open(
+            id=id,
+            name=name,
+            storage_dir=storage_dir,
+        )
+
+        if purge_on_start:
+            await dataset_client.drop()
+            dataset_client = await FileSystemDatasetClient.open(
+                id=id,
+                name=name,
+                storage_dir=storage_dir,
+            )
+
+        return dataset_client
+
+    @override
+    async def open_key_value_store_client(
+        self,
+        *,
+        id: str | None,
+        name: str | None,
+        purge_on_start: bool,
+        storage_dir: Path,
+    ) -> FileSystemKeyValueStoreClient:
+        return FileSystemKeyValueStoreClient()
+
+    @override
+    async def open_request_queue_client(
+        self,
+        *,
+        id: str | None,
+        name: str | None,
+        purge_on_start: bool,
+        storage_dir: Path,
+    ) -> FileSystemRequestQueueClient:
+        return FileSystemRequestQueueClient()

src/crawlee/storage_clients/_memory/__init__.py

@@ -1,3 +1,3 @@
-from ._storage_client import memory_storage_client
+from ._storage_client import MemoryStorageClient
 
-__all__ = ['memory_storage_client']
+__all__ = ['MemoryStorageClient']

src/crawlee/storage_clients/_memory/_storage_client.py

@@ -1,13 +1,57 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
+from typing_extensions import override
+
 from crawlee.storage_clients._base import StorageClient
 
 from ._dataset_client import MemoryDatasetClient
 from ._key_value_store_client import MemoryKeyValueStoreClient
 from ._request_queue_client import MemoryRequestQueueClient
 
-memory_storage_client = StorageClient(
-    dataset_client_class=MemoryDatasetClient,
-    key_value_store_client_class=MemoryKeyValueStoreClient,
-    request_queue_client_class=MemoryRequestQueueClient,
-)
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+class MemoryStorageClient(StorageClient):
+    """Memory storage client."""
+
+    @override
+    async def open_dataset_client(
+        self,
+        *,
+        id: str | None,
+        name: str | None,
+        purge_on_start: bool,
+        storage_dir: Path,
+    ) -> MemoryDatasetClient:
+        dataset_client = await MemoryDatasetClient.open(id=id, name=name, storage_dir=storage_dir)
+
+        if purge_on_start:
+            await dataset_client.drop()
+            dataset_client = await MemoryDatasetClient.open(id=id, name=name, storage_dir=storage_dir)
+
+        return dataset_client
+
+    @override
+    async def open_key_value_store_client(
+        self,
+        *,
+        id: str | None,
+        name: str | None,
+        purge_on_start: bool,
+        storage_dir: Path,
+    ) -> MemoryKeyValueStoreClient:
+        return MemoryKeyValueStoreClient()
+
+    @override
+    async def open_request_queue_client(
+        self,
+        *,
+        id: str | None,
+        name: str | None,
+        purge_on_start: bool,
+        storage_dir: Path,
+    ) -> MemoryRequestQueueClient:
+        return MemoryRequestQueueClient()

src/crawlee/storages/_dataset.py

@@ -31,6 +31,7 @@
 # TODO:
 # - inherit from storage class
 # - export methods
+# - caching / memoization of both datasets & dataset clients
 
 # Dataset
 # - properties:
@@ -116,30 +117,25 @@ async def open(
         id: str | None = None,
         name: str | None = None,
         purge_on_start: bool | None = None,
+        storage_dir: Path | None = None,
         configuration: Configuration | None = None,
         storage_client: StorageClient | None = None,
     ) -> Dataset:
         if id and name:
             raise ValueError('Only one of "id" or "name" can be specified, not both.')
 
-        configuration = configuration or service_locator.get_configuration()
-        storage_client = storage_client or service_locator.get_storage_client()
+        configuration = service_locator.get_configuration() if configuration is None else configuration
+        storage_client = service_locator.get_storage_client() if storage_client is None else storage_client
         purge_on_start = configuration.purge_on_start if purge_on_start is None else purge_on_start
+        storage_dir = Path(configuration.storage_dir) if storage_dir is None else storage_dir
 
-        dataset_client = await storage_client.dataset_client_class.open(
+        dataset_client = await storage_client.open_dataset_client(
             id=id,
             name=name,
-            storage_dir=Path(configuration.storage_dir),
+            purge_on_start=purge_on_start,
+            storage_dir=storage_dir,
         )
 
-        if purge_on_start:
-            await dataset_client.drop()
-            dataset_client = await storage_client.dataset_client_class.open(
-                id=id,
-                name=name,
-                storage_dir=Path(configuration.storage_dir),
-            )
-
         return cls(dataset_client)
 
     async def drop(self) -> None:

website/generate_module_shortcuts.py

@@ -5,6 +5,7 @@
 import importlib
 import inspect
 import json
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -55,5 +56,5 @@ def resolve_shortcuts(shortcuts: dict) -> None:
 
 resolve_shortcuts(shortcuts)
 
-with open('module_shortcuts.json', 'w', encoding='utf-8') as shortcuts_file:
+with Path('module_shortcuts.json').open('w', encoding='utf-8') as shortcuts_file:
     json.dump(shortcuts, shortcuts_file, indent=4, sort_keys=True)