diff --git a/changes/3359.feature.rst b/changes/3359.feature.rst new file mode 100644 index 0000000000..97a9a6e801 --- /dev/null +++ b/changes/3359.feature.rst @@ -0,0 +1 @@ +Expose `LatencyStore` as public API. \ No newline at end of file diff --git a/src/zarr/storage/__init__.py b/src/zarr/storage/__init__.py index 00df50214f..4b09a44cbb 100644 --- a/src/zarr/storage/__init__.py +++ b/src/zarr/storage/__init__.py @@ -6,6 +6,7 @@ from zarr.errors import ZarrDeprecationWarning from zarr.storage._common import StoreLike, StorePath from zarr.storage._fsspec import FsspecStore +from zarr.storage._latency import LatencyStore from zarr.storage._local import LocalStore from zarr.storage._logging import LoggingStore from zarr.storage._memory import GpuMemoryStore, MemoryStore @@ -16,6 +17,7 @@ __all__ = [ "FsspecStore", "GpuMemoryStore", + "LatencyStore", "LocalStore", "LoggingStore", "MemoryStore", diff --git a/src/zarr/storage/_latency.py b/src/zarr/storage/_latency.py new file mode 100644 index 0000000000..12ad9a85a5 --- /dev/null +++ b/src/zarr/storage/_latency.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING + +from zarr.abc.store import ByteRequest, Store +from zarr.storage._wrapper import WrapperStore + +if TYPE_CHECKING: + from zarr.core.buffer import Buffer, BufferPrototype + + +class LatencyStore(WrapperStore[Store]): + """ + A wrapper class that takes any store class in its constructor and + adds latency to the `set` and `get` methods. This can be used for + performance testing. + + Particularly useful for testing downstream applications which will + interact with a high-latency zarr store implementation, + such as one which read from or writes to remote object storage. + For example, by using this class to wrap a ``MemoryStore`` instance, + you can (crudely) simulate the latency of reading and writing from S3 + without having to actually use the network, or a mock like MinIO. + + Parameters + ---------- + store : Store + Store to wrap + get_latency : float + Amount of latency to add to each get call, in seconds. Default is 0. + set_latency : float + Amount of latency to add to each set call, in seconds. Default is 0. + """ + + get_latency: float + set_latency: float + + def __init__(self, cls: Store, *, get_latency: float = 0, set_latency: float = 0) -> None: + self.get_latency = float(get_latency) + self.set_latency = float(set_latency) + self._store = cls + + async def set(self, key: str, value: Buffer) -> None: + """ + Add latency to the ``set`` method. + + Calls ``asyncio.sleep(self.set_latency)`` before invoking the wrapped ``set`` method. + + Parameters + ---------- + key : str + The key to set + value : Buffer + The value to set + + Returns + ------- + None + """ + await asyncio.sleep(self.set_latency) + await self._store.set(key, value) + + async def get( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + """ + Add latency to the ``get`` method. + + Calls ``asyncio.sleep(self.get_latency)`` before invoking the wrapped ``get`` method. + + Parameters + ---------- + key : str + The key to get + prototype : BufferPrototype + The BufferPrototype to use. + byte_range : ByteRequest, optional + An optional byte range. + + Returns + ------- + buffer : Buffer or None + """ + await asyncio.sleep(self.get_latency) + return await self._store.get(key, prototype=prototype, byte_range=byte_range) diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index d2946705f0..2cca4118f4 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -1,17 +1,13 @@ from __future__ import annotations -import asyncio import pickle from abc import abstractmethod from typing import TYPE_CHECKING, Generic, TypeVar -from zarr.storage import WrapperStore - if TYPE_CHECKING: from typing import Any from zarr.abc.store import ByteRequest - from zarr.core.buffer.core import BufferPrototype import pytest @@ -525,63 +521,3 @@ async def test_set_if_not_exists(self, store: S) -> None: result = await store.get("k2", default_buffer_prototype()) assert result == new - - -class LatencyStore(WrapperStore[Store]): - """ - A wrapper class that takes any store class in its constructor and - adds latency to the `set` and `get` methods. This can be used for - performance testing. - """ - - get_latency: float - set_latency: float - - def __init__(self, cls: Store, *, get_latency: float = 0, set_latency: float = 0) -> None: - self.get_latency = float(get_latency) - self.set_latency = float(set_latency) - self._store = cls - - async def set(self, key: str, value: Buffer) -> None: - """ - Add latency to the ``set`` method. - - Calls ``asyncio.sleep(self.set_latency)`` before invoking the wrapped ``set`` method. - - Parameters - ---------- - key : str - The key to set - value : Buffer - The value to set - - Returns - ------- - None - """ - await asyncio.sleep(self.set_latency) - await self._store.set(key, value) - - async def get( - self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None - ) -> Buffer | None: - """ - Add latency to the ``get`` method. - - Calls ``asyncio.sleep(self.get_latency)`` before invoking the wrapped ``get`` method. - - Parameters - ---------- - key : str - The key to get - prototype : BufferPrototype - The BufferPrototype to use. - byte_range : ByteRequest, optional - An optional byte range. - - Returns - ------- - buffer : Buffer or None - """ - await asyncio.sleep(self.get_latency) - return await self._store.get(key, prototype=prototype, byte_range=byte_range) diff --git a/tests/test_group.py b/tests/test_group.py index e5cfe82daa..31ad14d62f 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -47,10 +47,9 @@ ZarrDeprecationWarning, ZarrUserWarning, ) -from zarr.storage import LocalStore, MemoryStore, StorePath, ZipStore +from zarr.storage import LatencyStore, LocalStore, MemoryStore, StorePath, ZipStore from zarr.storage._common import make_store_path from zarr.storage._utils import _join_paths, normalize_path -from zarr.testing.store import LatencyStore from .conftest import meta_from_array, parse_store