Expose LatencyStore as public API

changes/3359.feature.rst

@@ -0,0 +1 @@
+Expose `LatencyStore` as public API.

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",

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)

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)

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