Fix typing in test_wrapper

Author: dstansby

pyproject.toml

@@ -356,6 +356,7 @@ module = [
     "tests.test_store.test_logging",
     "tests.test_store.test_object",
     "tests.test_store.test_stateful",
+    "tests.test_store.test_wrapper",
 ]
 strict = false
 
@@ -364,7 +365,6 @@ strict = false
 [[tool.mypy.overrides]]
 module = [
     "tests.test_metadata.*",
-    "tests.test_store.test_wrapper",
     "tests.test_group",
     "tests.test_indexing",
     "tests.test_properties",

src/zarr/storage/_wrapper.py

@@ -7,8 +7,9 @@
     from types import TracebackType
     from typing import Any, Self
 
+    from zarr.abc.buffer import Buffer
     from zarr.abc.store import ByteRequest
-    from zarr.core.buffer import Buffer, BufferPrototype
+    from zarr.core.buffer import BufferPrototype
     from zarr.core.common import BytesLike
 
 from zarr.abc.store import Store

tests/test_store/test_wrapper.py

@@ -1,72 +1,82 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, TypedDict
 
 import pytest
 
-from zarr.core.buffer.cpu import Buffer, buffer_prototype
+from zarr.abc.store import ByteRequest, Store
+from zarr.core.buffer import Buffer
+from zarr.core.buffer.cpu import buffer_prototype
 from zarr.storage import LocalStore, WrapperStore
 from zarr.testing.store import StoreTests
 
 if TYPE_CHECKING:
-    from _pytest.compat import LEGACY_PATH
+    from pathlib import Path
 
-    from zarr.abc.store import Store
     from zarr.core.buffer.core import BufferPrototype
 
 
+class StoreKwargs(TypedDict):
+    store: LocalStore
+
+
+class OpenKwargs(TypedDict):
+    store_cls: type[LocalStore]
+    root: str
+
+
 # TODO: fix this warning
 @pytest.mark.filterwarnings(
     "ignore:coroutine 'ClientCreatorContext.__aexit__' was never awaited:RuntimeWarning"
 )
-class TestWrapperStore(StoreTests[WrapperStore, Buffer]):
-    store_cls = WrapperStore
+class TestWrapperStore(StoreTests[WrapperStore[LocalStore], Buffer]):
+    store_cls = WrapperStore[LocalStore]
     buffer_cls = Buffer
 
-    async def get(self, store: WrapperStore, key: str) -> Buffer:
+    async def get(self, store: WrapperStore[LocalStore], key: str) -> Buffer:
         return self.buffer_cls.from_bytes((store._store.root / key).read_bytes())
 
-    async def set(self, store: WrapperStore, key: str, value: Buffer) -> None:
+    async def set(self, store: WrapperStore[LocalStore], key: str, value: Buffer) -> None:
         parent = (store._store.root / key).parent
         if not parent.exists():
             parent.mkdir(parents=True)
         (store._store.root / key).write_bytes(value.to_bytes())
 
     @pytest.fixture
-    def store_kwargs(self, tmpdir: LEGACY_PATH) -> dict[str, str]:
-        return {"store": LocalStore(str(tmpdir))}
+    def store_kwargs(self, tmp_path: Path) -> StoreKwargs:
+        return {"store": LocalStore(str(tmp_path))}
 
     @pytest.fixture
-    def open_kwargs(self, tmpdir) -> dict[str, str]:
-        return {"store_cls": LocalStore, "root": str(tmpdir)}
+    def open_kwargs(self, tmp_path: Path) -> OpenKwargs:
+        return {"store_cls": LocalStore, "root": str(tmp_path)}
 
-    def test_store_supports_writes(self, store: WrapperStore) -> None:
+    def test_store_supports_writes(self, store: WrapperStore[LocalStore]) -> None:
         assert store.supports_writes
 
-    def test_store_supports_partial_writes(self, store: WrapperStore) -> None:
+    def test_store_supports_partial_writes(self, store: WrapperStore[LocalStore]) -> None:
         assert store.supports_partial_writes
 
-    def test_store_supports_listing(self, store: WrapperStore) -> None:
+    def test_store_supports_listing(self, store: WrapperStore[LocalStore]) -> None:
         assert store.supports_listing
 
-    def test_store_repr(self, store: WrapperStore) -> None:
+    def test_store_repr(self, store: WrapperStore[LocalStore]) -> None:
         assert f"{store!r}" == f"WrapperStore(LocalStore, 'file://{store._store.root.as_posix()}')"
 
-    def test_store_str(self, store: WrapperStore) -> None:
+    def test_store_str(self, store: WrapperStore[LocalStore]) -> None:
         assert str(store) == f"wrapping-file://{store._store.root.as_posix()}"
 
-    def test_check_writeable(self, store: WrapperStore) -> None:
+    def test_check_writeable(self, store: WrapperStore[LocalStore]) -> None:
         """
         Test _check_writeable() runs without errors.
         """
         store._check_writable()
 
-    def test_close(self, store: WrapperStore) -> None:
+    def test_close(self, store: WrapperStore[LocalStore]) -> None:
         "Test store can be closed"
         store.close()
         assert not store._is_open
 
-    def test_is_open_setter_raises(self, store: WrapperStore) -> None:
+    def test_is_open_setter_raises(self, store: WrapperStore[LocalStore]) -> None:
         """
         Test that a user cannot change `_is_open` without opening the underlying store.
         """
@@ -83,7 +93,7 @@ def test_is_open_setter_raises(self, store: WrapperStore) -> None:
 @pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=True)
 async def test_wrapped_set(store: Store, capsys: pytest.CaptureFixture[str]) -> None:
     # define a class that prints when it sets
-    class NoisySetter(WrapperStore):
+    class NoisySetter(WrapperStore[Store]):
         async def set(self, key: str, value: Buffer) -> None:
             print(f"setting {key}")
             await super().set(key, value)
@@ -101,15 +111,17 @@ async def set(self, key: str, value: Buffer) -> None:
 @pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=True)
 async def test_wrapped_get(store: Store, capsys: pytest.CaptureFixture[str]) -> None:
     # define a class that prints when it sets
-    class NoisyGetter(WrapperStore):
-        def get(self, key: str, prototype: BufferPrototype) -> None:
+    class NoisyGetter(WrapperStore[Any]):
+        async def get(
+            self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None
+        ) -> None:
             print(f"getting {key}")
-            return super().get(key, prototype=prototype)
+            await super().get(key, prototype=prototype, byte_range=byte_range)
 
     key = "foo"
     value = Buffer.from_bytes(b"bar")
     store_wrapped = NoisyGetter(store)
     await store_wrapped.set(key, value)
-    assert await store_wrapped.get(key, buffer_prototype) == value
+    await store_wrapped.get(key, buffer_prototype)
     captured = capsys.readouterr()
     assert f"getting {key}" in captured.out