use protocol

Author: d-v-b

src/zarr/storage/_fsspec.py

@@ -255,13 +255,19 @@ def from_url(
 
         return cls(fs=fs, path=path, read_only=read_only, allowed_exceptions=allowed_exceptions)
 
+    def get_semaphore(self) -> asyncio.Semaphore | None:
+        return self._semaphore
+
     def with_read_only(self, read_only: bool = False) -> FsspecStore:
         # docstring inherited
+        sem = self.get_semaphore()
+        concurrency_limit = sem._value if sem else None
         return type(self)(
             fs=self.fs,
             path=self.path,
             allowed_exceptions=self.allowed_exceptions,
             read_only=read_only,
+            concurrency_limit=concurrency_limit,
         )
 
     async def clear(self) -> None:
@@ -353,15 +359,16 @@ async def _set_many(self, values: Iterable[tuple[str, Buffer]]) -> None:
         if not self._is_open:
             await self._open()
         self._check_writable()
+        semaphore = self.get_semaphore()
 
         async def _set_with_limit(key: str, value: Buffer) -> None:
             if not isinstance(value, Buffer):
                 raise TypeError(
                     f"FsspecStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
                 )
             path = _dereference_path(self.path, key)
-            if self._semaphore:
-                async with self._semaphore:
+            if semaphore:
+                async with semaphore:
                     await self.fs._pipe_file(path, value.to_bytes())
             else:
                 await self.fs._pipe_file(path, value.to_bytes())

src/zarr/storage/_local.py

@@ -134,9 +134,13 @@ def __init__(
             asyncio.Semaphore(concurrency_limit) if concurrency_limit is not None else None
         )
 
+    def get_semaphore(self) -> asyncio.Semaphore | None:
+        return self._semaphore
+
     def with_read_only(self, read_only: bool = False) -> Self:
         # docstring inherited
-        concurrency_limit = self._semaphore._value if self._semaphore else None
+        sem = self.get_semaphore()
+        concurrency_limit = sem._value if sem else None
         return type(self)(
             root=self.root,
             read_only=read_only,
@@ -232,11 +236,13 @@ async def get_partial_values(
         # Note: We directly call the I/O functions here, wrapped with semaphore
         # to avoid deadlock from calling the decorated get() method
 
+        semaphore = self.get_semaphore()
+
         async def _get_with_limit(key: str, byte_range: ByteRequest | None) -> Buffer | None:
             path = self.root / key
             try:
-                if self._semaphore:
-                    async with self._semaphore:
+                if semaphore:
+                    async with semaphore:
                         return await asyncio.to_thread(_get, path, prototype, byte_range)
                 else:
                     return await asyncio.to_thread(_get, path, prototype, byte_range)

src/zarr/storage/_obstore.py

@@ -86,9 +86,13 @@ def __init__(
             asyncio.Semaphore(concurrency_limit) if concurrency_limit is not None else None
         )
 
+    def get_semaphore(self) -> asyncio.Semaphore | None:
+        return self._semaphore
+
     def with_read_only(self, read_only: bool = False) -> Self:
         # docstring inherited
-        concurrency_limit = self._semaphore._value if self._semaphore else None
+        sem = self.get_semaphore()
+        concurrency_limit = sem._value if sem else None
         return type(self)(
             store=self.store,
             read_only=read_only,
@@ -134,6 +138,7 @@ async def get_partial_values(
         import obstore as obs
 
         key_ranges = list(key_ranges)
+        semaphore = self.get_semaphore()
         # Group bounded range requests by path for batched fetching
         per_file_bounded: dict[str, list[tuple[int, RangeByteRequest]]] = defaultdict(list)
         other_requests: list[tuple[int, str, ByteRequest | None]] = []
@@ -150,8 +155,8 @@ async def _fetch_ranges(path: str, requests: list[tuple[int, RangeByteRequest]])
             """Batch multiple range requests for the same file using get_ranges_async."""
             starts = [r.start for _, r in requests]
             ends = [r.end for _, r in requests]
-            if self._semaphore:
-                async with self._semaphore:
+            if semaphore:
+                async with semaphore:
                     responses = await obs.get_ranges_async(
                         self.store, path=path, starts=starts, ends=ends
                     )
@@ -165,8 +170,8 @@ async def _fetch_ranges(path: str, requests: list[tuple[int, RangeByteRequest]])
         async def _fetch_one(idx: int, path: str, byte_range: ByteRequest | None) -> None:
             """Fetch a single non-range request with semaphore limiting."""
             try:
-                if self._semaphore:
-                    async with self._semaphore:
+                if semaphore:
+                    async with semaphore:
                         buffers[idx] = await self._get_impl(path, prototype, byte_range, obs)
                 else:
                     buffers[idx] = await self._get_impl(path, prototype, byte_range, obs)
@@ -250,11 +255,12 @@ async def _set_many(self, values: Iterable[tuple[str, Buffer]]) -> None:
         import obstore as obs
 
         self._check_writable()
+        semaphore = self.get_semaphore()
 
         async def _set_with_limit(key: str, value: Buffer) -> None:
             buf = value.as_buffer_like()
-            if self._semaphore:
-                async with self._semaphore:
+            if semaphore:
+                async with semaphore:
                     await obs.put_async(self.store, key, buf)
             else:
                 await obs.put_async(self.store, key, buf)
@@ -268,8 +274,9 @@ async def set_if_not_exists(self, key: str, value: Buffer) -> None:
 
         self._check_writable()
         buf = value.as_buffer_like()
-        if self._semaphore:
-            async with self._semaphore:
+        semaphore = self.get_semaphore()
+        if semaphore:
+            async with semaphore:
                 with contextlib.suppress(obs.exceptions.AlreadyExistsError):
                     await obs.put_async(self.store, key, buf, mode="create")
         else:
@@ -304,11 +311,12 @@ async def delete_dir(self, prefix: str) -> None:
             prefix += "/"
 
         metas = await obs.list(self.store, prefix).collect_async()
+        semaphore = self.get_semaphore()
 
         # Delete with semaphore limiting to avoid deadlock
         async def _delete_with_limit(path: str) -> None:
-            if self._semaphore:
-                async with self._semaphore:
+            if semaphore:
+                async with semaphore:
                     with contextlib.suppress(FileNotFoundError):
                         await obs.delete_async(self.store, path)
             else:

src/zarr/storage/_utils.py

@@ -3,7 +3,7 @@
 import functools
 import re
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar
+from typing import TYPE_CHECKING, Any, ParamSpec, Protocol, TypeVar, runtime_checkable
 
 from zarr.abc.store import OffsetByteRequest, RangeByteRequest, SuffixByteRequest
 
@@ -18,20 +18,24 @@
 T_co = TypeVar("T_co", covariant=True)
 
 
-def with_concurrency_limit(
-    semaphore_attr: str = "_semaphore",
-) -> Callable[[Callable[P, Coroutine[Any, Any, T_co]]], Callable[P, Coroutine[Any, Any, T_co]]]:
+@runtime_checkable
+class HasConcurrencyLimit(Protocol):
+    """Protocol for stores that support concurrency limiting via a semaphore."""
+
+    def get_semaphore(self) -> asyncio.Semaphore | None:
+        """Return the semaphore used for concurrency limiting, or None for unlimited."""
+        ...
+
+
+def with_concurrency_limit() -> Callable[
+    [Callable[P, Coroutine[Any, Any, T_co]]], Callable[P, Coroutine[Any, Any, T_co]]
+]:
     """
     Decorator that applies a semaphore-based concurrency limit to an async method.
 
-    This decorator is designed for Store methods that need to limit concurrent operations.
-    The store instance should have a `_semaphore` attribute (or custom attribute name)
-    that is either an asyncio.Semaphore or None (for unlimited concurrency).
-
-    Parameters
-    ----------
-    semaphore_attr : str, optional
-        Name of the semaphore attribute on the class instance. Default is "_semaphore".
+    This decorator is designed for methods on classes that implement the
+    ``HasConcurrencyLimit`` protocol. The class must define a ``get_semaphore()``
+    method returning either an ``asyncio.Semaphore`` or ``None``.
 
     Returns
     -------
@@ -45,6 +49,9 @@ class MyStore(Store):
         def __init__(self, concurrency_limit: int = 100):
             self._semaphore = asyncio.Semaphore(concurrency_limit) if concurrency_limit else None
 
+        def get_semaphore(self) -> asyncio.Semaphore | None:
+            return self._semaphore
+
         @with_concurrency_limit()
         async def get(self, key: str) -> Buffer | None:
             # This will only run when semaphore permits
@@ -55,22 +62,14 @@ async def get(self, key: str) -> Buffer | None:
     def decorator(
         func: Callable[P, Coroutine[Any, Any, T_co]],
     ) -> Callable[P, Coroutine[Any, Any, T_co]]:
-        """
-        This decorator wraps the invocation of `func` in an `async with semaphore` context manager.
-        The semaphore object is resolved by getting the `semaphor_attr` attribute from the first
-        argument to func. When this decorator is used on a method of a class, that first argument
-        is a reference to the class instance (`self`).
-        """
-
         @functools.wraps(func)
         async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T_co:
             # First arg should be 'self'
             if not args:
                 raise TypeError(f"{func.__name__} requires at least one argument (self)")
 
             self = args[0]
-
-            semaphore: asyncio.Semaphore | None = getattr(self, semaphore_attr)
+            semaphore: asyncio.Semaphore | None = self.get_semaphore()  # type: ignore[attr-defined]
 
             if semaphore is not None:
                 async with semaphore:

src/zarr/testing/store_concurrency.py

@@ -40,16 +40,25 @@ async def store(self, store_kwargs: dict[str, Any]) -> S:
         """Create and open a store instance."""
         return await self.store_cls.open(**store_kwargs)
 
+    @staticmethod
+    def _get_semaphore(store: Store) -> asyncio.Semaphore | None:
+        """Get the semaphore from a store, or None if the store doesn't support concurrency limiting."""
+        get_semaphore = getattr(store, "get_semaphore", None)
+        if get_semaphore is not None:
+            return get_semaphore()  # type: ignore[no-any-return]
+        return None
+
     def test_concurrency_limit_default(self, store: S) -> None:
         """Test that store has the expected default concurrency limit."""
-        if hasattr(store, "_semaphore"):
-            if self.expected_concurrency_limit is None:
-                assert store._semaphore is None, "Expected no concurrency limit"
-            else:
-                assert store._semaphore is not None, "Expected concurrency limit to be set"
-                assert store._semaphore._value == self.expected_concurrency_limit, (
-                    f"Expected limit {self.expected_concurrency_limit}, got {store._semaphore._value}"
-                )
+        semaphore = self._get_semaphore(store)
+        if semaphore is None and self.expected_concurrency_limit is not None:
+            pytest.fail("Expected concurrency limit to be set")
+        if semaphore is not None and self.expected_concurrency_limit is None:
+            pytest.fail("Expected no concurrency limit")
+        if semaphore is not None and self.expected_concurrency_limit is not None:
+            assert semaphore._value == self.expected_concurrency_limit, (
+                f"Expected limit {self.expected_concurrency_limit}, got {semaphore._value}"
+            )
 
     def test_concurrency_limit_custom(self, store_kwargs: dict[str, Any]) -> None:
         """Test that custom concurrency limits can be set."""
@@ -58,25 +67,25 @@ def test_concurrency_limit_custom(self, store_kwargs: dict[str, Any]) -> None:
 
         # Test with custom limit
         store = self.store_cls(**{**store_kwargs, "concurrency_limit": 42})
-        if hasattr(store, "_semaphore"):
-            assert store._semaphore is not None
-            assert store._semaphore._value == 42
+        semaphore = self._get_semaphore(store)
+        assert semaphore is not None
+        assert semaphore._value == 42
 
         # Test with None (unlimited)
         store = self.store_cls(**{**store_kwargs, "concurrency_limit": None})
-        if hasattr(store, "_semaphore"):
-            assert store._semaphore is None
+        assert self._get_semaphore(store) is None
 
     async def test_concurrency_limit_enforced(self, store: S) -> None:
         """Test that the concurrency limit is actually enforced during execution.
 
         This test verifies that when many operations are submitted concurrently,
         only up to the concurrency limit are actually executing at once.
         """
-        if not hasattr(store, "_semaphore") or store._semaphore is None:
+        semaphore = self._get_semaphore(store)
+        if semaphore is None:
             pytest.skip("Store has no concurrency limit")
 
-        limit = store._semaphore._value
+        limit = semaphore._value
 
         # We'll monitor the semaphore's available count
         # When it reaches 0, that means `limit` operations are running
@@ -86,7 +95,7 @@ async def monitored_operation(key: str, value: B) -> None:
             nonlocal min_available
             # Check semaphore state right after we're scheduled
             await asyncio.sleep(0)  # Yield to ensure we're in the queue
-            available = store._semaphore._value
+            available = semaphore._value
             min_available = min(min_available, available)
 
             # Now do the actual operation (which will acquire the semaphore)