use base class instead of protocol

Author: d-v-b

src/zarr/storage/__init__.py

@@ -10,10 +10,12 @@
 from zarr.storage._logging import LoggingStore
 from zarr.storage._memory import GpuMemoryStore, MemoryStore
 from zarr.storage._obstore import ObjectStore
+from zarr.storage._utils import ConcurrencyLimiter
 from zarr.storage._wrapper import WrapperStore
 from zarr.storage._zip import ZipStore
 
 __all__ = [
+    "ConcurrencyLimiter",
     "FsspecStore",
     "GpuMemoryStore",
     "LocalStore",

src/zarr/storage/_fsspec.py

@@ -18,7 +18,7 @@
 from zarr.core.buffer import Buffer
 from zarr.errors import ZarrUserWarning
 from zarr.storage._common import _dereference_path
-from zarr.storage._utils import with_concurrency_limit
+from zarr.storage._utils import ConcurrencyLimiter, with_concurrency_limit
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterable
@@ -69,7 +69,7 @@ def _make_async(fs: AbstractFileSystem) -> AsyncFileSystem:
     return AsyncFileSystemWrapper(fs, asynchronous=True)
 
 
-class FsspecStore(Store):
+class FsspecStore(Store, ConcurrencyLimiter):
     """
     Store for remote data based on FSSpec.
 
@@ -122,7 +122,6 @@ class FsspecStore(Store):
     fs: AsyncFileSystem
     allowed_exceptions: tuple[type[Exception], ...]
     path: str
-    _semaphore: asyncio.Semaphore | None
 
     def __init__(
         self,
@@ -133,13 +132,11 @@ def __init__(
         allowed_exceptions: tuple[type[Exception], ...] = ALLOWED_EXCEPTIONS,
         concurrency_limit: int | None = 50,
     ) -> None:
-        super().__init__(read_only=read_only)
+        Store.__init__(self, read_only=read_only)
+        ConcurrencyLimiter.__init__(self, concurrency_limit)
         self.fs = fs
         self.path = path
         self.allowed_exceptions = allowed_exceptions
-        self._semaphore = (
-            asyncio.Semaphore(concurrency_limit) if concurrency_limit is not None else None
-        )
 
         if not self.fs.async_impl:
             raise TypeError("Filesystem needs to support async operations.")
@@ -255,19 +252,14 @@ 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,
+            concurrency_limit=self.concurrency_limit,
         )
 
     async def clear(self) -> None:
@@ -290,7 +282,7 @@ def __eq__(self, other: object) -> bool:
             and self.fs == other.fs
         )
 
-    @with_concurrency_limit()
+    @with_concurrency_limit
     async def get(
         self,
         key: str,
@@ -333,7 +325,7 @@ async def get(
         else:
             return value
 
-    @with_concurrency_limit()
+    @with_concurrency_limit
     async def set(
         self,
         key: str,
@@ -359,23 +351,19 @@ 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 semaphore:
-                async with semaphore:
-                    await self.fs._pipe_file(path, value.to_bytes())
-            else:
+            async with self._limit():
                 await self.fs._pipe_file(path, value.to_bytes())
 
         await asyncio.gather(*[_set_with_limit(key, value) for key, value in values])
 
-    @with_concurrency_limit()
+    @with_concurrency_limit
     async def delete(self, key: str) -> None:
         # docstring inherited
         self._check_writable()

src/zarr/storage/_local.py

@@ -19,7 +19,7 @@
 )
 from zarr.core.buffer import Buffer
 from zarr.core.buffer.core import default_buffer_prototype
-from zarr.storage._utils import with_concurrency_limit
+from zarr.storage._utils import ConcurrencyLimiter, with_concurrency_limit
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterable, Iterator
@@ -86,7 +86,7 @@ def _put(path: Path, value: Buffer, exclusive: bool = False) -> int:
         return f.write(view)
 
 
-class LocalStore(Store):
+class LocalStore(Store, ConcurrencyLimiter):
     """
     Store for the local file system.
 
@@ -113,7 +113,6 @@ class LocalStore(Store):
     supports_listing: bool = True
 
     root: Path
-    _semaphore: asyncio.Semaphore | None
 
     def __init__(
         self,
@@ -122,29 +121,22 @@ def __init__(
         read_only: bool = False,
         concurrency_limit: int | None = 100,
     ) -> None:
-        super().__init__(read_only=read_only)
         if isinstance(root, str):
             root = Path(root)
         if not isinstance(root, Path):
             raise TypeError(
                 f"'root' must be a string or Path instance. Got an instance of {type(root)} instead."
             )
+        Store.__init__(self, read_only=read_only)
+        ConcurrencyLimiter.__init__(self, concurrency_limit)
         self.root = root
-        self._semaphore = (
-            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
-        sem = self.get_semaphore()
-        concurrency_limit = sem._value if sem else None
         return type(self)(
             root=self.root,
             read_only=read_only,
-            concurrency_limit=concurrency_limit,
+            concurrency_limit=self.concurrency_limit,
         )
 
     @classmethod
@@ -207,7 +199,7 @@ def __repr__(self) -> str:
     def __eq__(self, other: object) -> bool:
         return isinstance(other, type(self)) and self.root == other.root
 
-    @with_concurrency_limit()
+    @with_concurrency_limit
     async def get(
         self,
         key: str,
@@ -233,18 +225,13 @@ async def get_partial_values(
         key_ranges: Iterable[tuple[str, ByteRequest | None]],
     ) -> list[Buffer | None]:
         # docstring inherited
-        # 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()
+        # We directly call the I/O functions here, wrapped with the semaphore,
+        # to avoid deadlock from calling the decorated get() method.
 
         async def _get_with_limit(key: str, byte_range: ByteRequest | None) -> Buffer | None:
             path = self.root / key
             try:
-                if semaphore:
-                    async with semaphore:
-                        return await asyncio.to_thread(_get, path, prototype, byte_range)
-                else:
+                async with self._limit():
                     return await asyncio.to_thread(_get, path, prototype, byte_range)
             except (FileNotFoundError, IsADirectoryError, NotADirectoryError):
                 return None
@@ -264,7 +251,7 @@ async def set_if_not_exists(self, key: str, value: Buffer) -> None:
         except FileExistsError:
             pass
 
-    @with_concurrency_limit()
+    @with_concurrency_limit
     async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None:
         if not self._is_open:
             await self._open()
@@ -277,7 +264,7 @@ async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None:
         path = self.root / key
         await asyncio.to_thread(_put, path, value, exclusive=exclusive)
 
-    @with_concurrency_limit()
+    @with_concurrency_limit
     async def delete(self, key: str) -> None:
         """
         Remove a key from the store.

src/zarr/storage/_obstore.py

@@ -15,7 +15,7 @@
     Store,
     SuffixByteRequest,
 )
-from zarr.storage._utils import _relativize_path, with_concurrency_limit
+from zarr.storage._utils import ConcurrencyLimiter, _relativize_path, with_concurrency_limit
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Coroutine, Iterable, Sequence
@@ -38,7 +38,7 @@
 T_Store = TypeVar("T_Store", bound="_UpstreamObjectStore")
 
 
-class ObjectStore(Store, Generic[T_Store]):
+class ObjectStore(Store, ConcurrencyLimiter, Generic[T_Store]):
     """
     Store that uses obstore for fast read/write from AWS, GCP, Azure.
 
@@ -60,7 +60,6 @@ class ObjectStore(Store, Generic[T_Store]):
 
     store: T_Store
     """The underlying obstore instance."""
-    _semaphore: asyncio.Semaphore | None
 
     def __eq__(self, value: object) -> bool:
         if not isinstance(value, ObjectStore):
@@ -80,23 +79,16 @@ def __init__(
     ) -> None:
         if not store.__class__.__module__.startswith("obstore"):
             raise TypeError(f"expected ObjectStore class, got {store!r}")
-        super().__init__(read_only=read_only)
+        Store.__init__(self, read_only=read_only)
+        ConcurrencyLimiter.__init__(self, concurrency_limit)
         self.store = store
-        self._semaphore = (
-            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
-        sem = self.get_semaphore()
-        concurrency_limit = sem._value if sem else None
         return type(self)(
             store=self.store,
             read_only=read_only,
-            concurrency_limit=concurrency_limit,
+            concurrency_limit=self.concurrency_limit,
         )
 
     def __str__(self) -> str:
@@ -114,7 +106,7 @@ def __setstate__(self, state: dict[Any, Any]) -> None:
         state["store"] = pickle.loads(state["store"])
         self.__dict__.update(state)
 
-    @with_concurrency_limit()
+    @with_concurrency_limit
     async def get(
         self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None
     ) -> Buffer | None:
@@ -138,7 +130,6 @@ 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]] = []
@@ -155,12 +146,7 @@ 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 semaphore:
-                async with semaphore:
-                    responses = await obs.get_ranges_async(
-                        self.store, path=path, starts=starts, ends=ends
-                    )
-            else:
+            async with self._limit():
                 responses = await obs.get_ranges_async(
                     self.store, path=path, starts=starts, ends=ends
                 )
@@ -170,10 +156,7 @@ 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 semaphore:
-                    async with semaphore:
-                        buffers[idx] = await self._get_impl(path, prototype, byte_range, obs)
-                else:
+                async with self._limit():
                     buffers[idx] = await self._get_impl(path, prototype, byte_range, obs)
             except _ALLOWED_EXCEPTIONS:
                 pass  # buffers[idx] stays None
@@ -240,7 +223,7 @@ def supports_writes(self) -> bool:
         # docstring inherited
         return True
 
-    @with_concurrency_limit()
+    @with_concurrency_limit
     async def set(self, key: str, value: Buffer) -> None:
         # docstring inherited
         import obstore as obs
@@ -255,31 +238,22 @@ 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 semaphore:
-                async with semaphore:
-                    await obs.put_async(self.store, key, buf)
-            else:
+            async with self._limit():
                 await obs.put_async(self.store, key, buf)
 
         await asyncio.gather(*[_set_with_limit(key, value) for key, value in values])
 
     async def set_if_not_exists(self, key: str, value: Buffer) -> None:
         # docstring inherited
-        # Note: Not decorated to avoid deadlock when called in batch via gather()
+        # Not decorated to avoid deadlock when called in batch via gather()
         import obstore as obs
 
         self._check_writable()
         buf = value.as_buffer_like()
-        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:
+        async with self._limit():
             with contextlib.suppress(obs.exceptions.AlreadyExistsError):
                 await obs.put_async(self.store, key, buf, mode="create")
 
@@ -288,7 +262,7 @@ def supports_deletes(self) -> bool:
         # docstring inherited
         return True
 
-    @with_concurrency_limit()
+    @with_concurrency_limit
     async def delete(self, key: str) -> None:
         # docstring inherited
         import obstore as obs
@@ -311,15 +285,9 @@ 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 semaphore:
-                async with semaphore:
-                    with contextlib.suppress(FileNotFoundError):
-                        await obs.delete_async(self.store, path)
-            else:
+            async with self._limit():
                 with contextlib.suppress(FileNotFoundError):
                     await obs.delete_async(self.store, path)