rm code duplication for open methods

Author: vdusek

src/crawlee/storages/_dataset.py

@@ -12,6 +12,7 @@
 
 from ._base import Storage
 from ._key_value_store import KeyValueStore
+from ._utils import open_storage_instance
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
@@ -109,39 +110,18 @@ async def open(
         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.')
-
-        # Check for default instance if no id or name provided
-        if id is None and name is None and cls._default_instance is not None:
-            return cls._default_instance
-
-        # Check if the dataset is already cached
-        if id is not None and id in cls._cache_by_id:
-            return cls._cache_by_id[id]
-        if name is not None and name in cls._cache_by_name:
-            return cls._cache_by_name[name]
-
         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
-
-        client = await storage_client.open_dataset_client(
+        return await open_storage_instance(
+            cls,
             id=id,
             name=name,
             configuration=configuration,
+            cache_by_id=cls._cache_by_id,
+            cache_by_name=cls._cache_by_name,
+            default_instance_attr='_default_instance',
+            client_opener=storage_client.open_dataset_client,
         )
-        dataset = cls(client)
-
-        # Cache the dataset instance by ID and name
-        cls._cache_by_id[dataset.id] = dataset
-        if dataset.name is not None:
-            cls._cache_by_name[dataset.name] = dataset
-
-        # Store as default instance if neither id nor name was provided
-        if id is None and name is None:
-            cls._default_instance = dataset
-
-        return dataset
 
     @override
     async def drop(self) -> None:

src/crawlee/storages/_key_value_store.py

@@ -14,6 +14,7 @@
 from crawlee.storage_clients.models import KeyValueStoreMetadata
 
 from ._base import Storage
+from ._utils import open_storage_instance
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
@@ -117,41 +118,19 @@ async def open(
         configuration: Configuration | None = None,
         storage_client: StorageClient | None = None,
     ) -> KeyValueStore:
-        if id and name:
-            raise ValueError('Only one of "id" or "name" can be specified, not both.')
-
-        # Check for default instance if no id or name provided
-        if id is None and name is None and cls._default_instance is not None:
-            return cls._default_instance
-
-        # Check if the key-value store is already cached
-        if id is not None and id in cls._cache_by_id:
-            return cls._cache_by_id[id]
-        if name is not None and name in cls._cache_by_name:
-            return cls._cache_by_name[name]
-
         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
-
-        client = await storage_client.open_key_value_store_client(
+        return await open_storage_instance(
+            cls,
             id=id,
             name=name,
             configuration=configuration,
+            cache_by_id=cls._cache_by_id,
+            cache_by_name=cls._cache_by_name,
+            default_instance_attr='_default_instance',
+            client_opener=storage_client.open_key_value_store_client,
         )
 
-        kvs = cls(client)
-
-        # Cache the key-value store instance by ID and name
-        cls._cache_by_id[kvs.id] = kvs
-        if kvs.name is not None:
-            cls._cache_by_name[kvs.name] = kvs
-
-        # Store as default instance if neither id nor name was provided
-        if id is None and name is None:
-            cls._default_instance = kvs
-
-        return kvs
-
     @override
     async def drop(self) -> None:
         if self.id in self._cache_by_id:

src/crawlee/storages/_request_queue.py

@@ -13,6 +13,7 @@
 from crawlee.request_loaders import RequestManager
 
 from ._base import Storage
+from ._utils import open_storage_instance
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -127,41 +128,19 @@ async def open(
         configuration: Configuration | None = None,
         storage_client: StorageClient | None = None,
     ) -> RequestQueue:
-        if id and name:
-            raise ValueError('Only one of "id" or "name" can be specified, not both.')
-
-        # Check for default instance if no id or name provided
-        if id is None and name is None and cls._default_instance is not None:
-            return cls._default_instance
-
-        # Check if the request queue is already cached
-        if id is not None and id in cls._cache_by_id:
-            return cls._cache_by_id[id]
-        if name is not None and name in cls._cache_by_name:
-            return cls._cache_by_name[name]
-
         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
-
-        client = await storage_client.open_request_queue_client(
+        return await open_storage_instance(
+            cls,
             id=id,
             name=name,
             configuration=configuration,
+            cache_by_id=cls._cache_by_id,
+            cache_by_name=cls._cache_by_name,
+            default_instance_attr='_default_instance',
+            client_opener=storage_client.open_request_queue_client,
         )
 
-        request_queue = cls(client)
-
-        # Cache the request queue instance by ID and name
-        cls._cache_by_id[request_queue.id] = request_queue
-        if request_queue.name is not None:
-            cls._cache_by_name[request_queue.name] = request_queue
-
-        # Store as default instance if neither id nor name was provided
-        if id is None and name is None:
-            cls._default_instance = request_queue
-
-        return request_queue
-
     @override
     async def drop(self) -> None:
         # Remove from cache before dropping

src/crawlee/storages/_utils.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from typing import Any, Callable, TypeVar, cast
+
+from ._base import Storage
+
+T = TypeVar('T', bound=Storage)
+
+
+async def open_storage_instance(
+    cls: type[T],
+    *,
+    id: str | None,
+    name: str | None,
+    configuration: Any,
+    cache_by_id: dict[str, T],
+    cache_by_name: dict[str, T],
+    default_instance_attr: str,
+    client_opener: Callable[..., Any],
+) -> T:
+    if id and name:
+        raise ValueError('Only one of "id" or "name" can be specified, not both.')
+
+    default_instance = getattr(cls, default_instance_attr)
+    if id is None and name is None and default_instance is not None:
+        return cast('T', default_instance)
+
+    if id is not None and id in cache_by_id:
+        return cache_by_id[id]
+    if name is not None and name in cache_by_name:
+        return cache_by_name[name]
+
+    client = await client_opener(id=id, name=name, configuration=configuration)
+    instance = cls(client)  # type: ignore[call-arg]
+    instance_name = getattr(instance, 'name', None)
+
+    cache_by_id[instance.id] = instance
+    if instance_name is not None:
+        cache_by_name[instance_name] = instance
+
+    if id is None and name is None:
+        setattr(cls, default_instance_attr, instance)
+
+    return instance