Fix typing in store tests (#3394)

* Fix typing in test_core

* Fix typing in test_logging

* Fix typing in test_object

* Fix typing in test_stateful

* Fix typing in test_wrapper

* Fix wrapper store tests

* Fix store in test

* Fix merge conflicts

* Adapt ObjectStore tests for new generic class

* Fix generic assignment in tests

---------

Co-authored-by: Davis Bennett <[email protected]>

Author: dstansby

pyproject.toml

@@ -359,18 +359,18 @@ module = [
     "tests.test_store.test_memory",
     "tests.test_codecs.test_codecs",
     "tests.test_metadata.*",
+    "tests.test_store.test_core",
+    "tests.test_store.test_logging",
+    "tests.test_store.test_object",
+    "tests.test_store.test_stateful",
+    "tests.test_store.test_wrapper",
 ]
 strict = false
 
 # TODO: Move the next modules up to the strict = false section
 # and fix the errors
 [[tool.mypy.overrides]]
 module = [
-    "tests.test_store.test_core",
-    "tests.test_store.test_logging",
-    "tests.test_store.test_object",
-    "tests.test_store.test_stateful",
-    "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.abc.store import Store
 

tests/test_store/test_core.py

@@ -1,5 +1,7 @@
 import tempfile
+from collections.abc import Callable, Generator
 from pathlib import Path
+from typing import Any, Literal
 
 import pytest
 from _pytest.compat import LEGACY_PATH
@@ -21,7 +23,9 @@
 @pytest.fixture(
     params=["none", "temp_dir_str", "temp_dir_path", "store_path", "memory_store", "dict"]
 )
-def store_like(request):
+def store_like(
+    request: pytest.FixtureRequest,
+) -> Generator[None | str | Path | StorePath | MemoryStore | dict[Any, Any], None, None]:
     if request.param == "none":
         yield None
     elif request.param == "temp_dir_str":
@@ -42,7 +46,7 @@ def store_like(request):
 @pytest.mark.parametrize("write_group", [True, False])
 @pytest.mark.parametrize("zarr_format", [2, 3])
 async def test_contains_group(
-    local_store, path: str, write_group: bool, zarr_format: ZarrFormat
+    local_store: LocalStore, path: str, write_group: bool, zarr_format: ZarrFormat
 ) -> None:
     """
     Test that the contains_group method correctly reports the existence of a group.
@@ -58,7 +62,7 @@ async def test_contains_group(
 @pytest.mark.parametrize("write_array", [True, False])
 @pytest.mark.parametrize("zarr_format", [2, 3])
 async def test_contains_array(
-    local_store, path: str, write_array: bool, zarr_format: ZarrFormat
+    local_store: LocalStore, path: str, write_array: bool, zarr_format: ZarrFormat
 ) -> None:
     """
     Test that the contains array method correctly reports the existence of an array.
@@ -71,13 +75,15 @@ async def test_contains_array(
 
 
 @pytest.mark.parametrize("func", [contains_array, contains_group])
-async def test_contains_invalid_format_raises(local_store, func: callable) -> None:
+async def test_contains_invalid_format_raises(
+    local_store: LocalStore, func: Callable[[Any], Any]
+) -> None:
     """
     Test contains_group and contains_array raise errors for invalid zarr_formats
     """
     store_path = StorePath(local_store)
     with pytest.raises(ValueError):
-        assert await func(store_path, zarr_format="3.0")
+        assert await func(store_path, zarr_format="3.0")  # type: ignore[call-arg]
 
 
 @pytest.mark.parametrize("path", [None, "", "bar"])
@@ -113,40 +119,48 @@ async def test_make_store_path_local(
 @pytest.mark.parametrize("path", [None, "", "bar"])
 @pytest.mark.parametrize("mode", ["r", "w"])
 async def test_make_store_path_store_path(
-    tmpdir: LEGACY_PATH, path: str, mode: AccessModeLiteral
+    tmp_path: Path, path: str, mode: AccessModeLiteral
 ) -> None:
     """
     Test invoking make_store_path when the input is another store_path. In particular we want to ensure
     that a new path is handled correctly.
     """
     ro = mode == "r"
-    store_like = await StorePath.open(LocalStore(str(tmpdir), read_only=ro), path="root", mode=mode)
+    store_like = await StorePath.open(
+        LocalStore(str(tmp_path), read_only=ro), path="root", mode=mode
+    )
     store_path = await make_store_path(store_like, path=path, mode=mode)
     assert isinstance(store_path.store, LocalStore)
-    assert Path(store_path.store.root) == Path(tmpdir)
+    assert Path(store_path.store.root) == tmp_path
     path_normalized = normalize_path(path)
     assert store_path.path == (store_like / path_normalized).path
     assert store_path.read_only == ro
 
 
 @pytest.mark.parametrize("modes", [(True, "w"), (False, "x")])
-async def test_store_path_invalid_mode_raises(tmpdir: LEGACY_PATH, modes: tuple) -> None:
+async def test_store_path_invalid_mode_raises(
+    tmp_path: Path, modes: tuple[bool, Literal["w", "x"]]
+) -> None:
     """
     Test that ValueErrors are raise for invalid mode.
     """
     with pytest.raises(ValueError):
-        await StorePath.open(LocalStore(str(tmpdir), read_only=modes[0]), path=None, mode=modes[1])
+        await StorePath.open(
+            LocalStore(str(tmp_path), read_only=modes[0]),
+            path="",
+            mode=modes[1],  # type:ignore[arg-type]
+        )
 
 
 async def test_make_store_path_invalid() -> None:
     """
     Test that invalid types raise TypeError
     """
     with pytest.raises(TypeError):
-        await make_store_path(1)  # type: ignore[arg-type]
+        await make_store_path(1)
 
 
-async def test_make_store_path_fsspec(monkeypatch) -> None:
+async def test_make_store_path_fsspec() -> None:
     pytest.importorskip("fsspec")
     pytest.importorskip("requests")
     pytest.importorskip("aiohttp")
@@ -161,7 +175,7 @@ async def test_make_store_path_storage_options_raises(store_like: StoreLike) ->
 
 async def test_unsupported() -> None:
     with pytest.raises(TypeError, match="Unsupported type for store_like: 'int'"):
-        await make_store_path(1)  # type: ignore[arg-type]
+        await make_store_path(1)
 
 
 @pytest.mark.parametrize(
@@ -184,12 +198,12 @@ def test_normalize_path_upath() -> None:
     assert normalize_path(upath.UPath("foo/bar")) == "foo/bar"
 
 
-def test_normalize_path_none():
+def test_normalize_path_none() -> None:
     assert normalize_path(None) == ""
 
 
 @pytest.mark.parametrize("path", [".", ".."])
-def test_normalize_path_invalid(path: str):
+def test_normalize_path_invalid(path: str) -> None:
     with pytest.raises(ValueError):
         normalize_path(path)
 
@@ -230,7 +244,7 @@ def test_invalid(paths: tuple[str, str]) -> None:
             _normalize_paths(paths)
 
 
-def test_normalize_path_keys():
+def test_normalize_path_keys() -> None:
     """
     Test that ``_normalize_path_keys`` just applies the normalize_path function to each key of its
     input
@@ -272,10 +286,10 @@ def test_different_open_mode(tmp_path: LEGACY_PATH) -> None:
 
     # Test with a store that doesn't implement .with_read_only()
     zarr_path = tmp_path / "foo.zarr"
-    store = ZipStore(zarr_path, mode="w")
-    zarr.create((100,), store=store, zarr_format=2, path="a")
+    zip_store = ZipStore(zarr_path, mode="w")
+    zarr.create((100,), store=zip_store, zarr_format=2, path="a")
     with pytest.raises(
         ValueError,
         match="Store is not read-only but mode is 'r'. Unable to create a read-only copy of the store. Please use a read-only store or a storage class that implements .with_read_only().",
     ):
-        zarr.open_array(store=store, path="a", zarr_format=2, mode="r")
+        zarr.open_array(store=zip_store, path="a", zarr_format=2, mode="r")

tests/test_store/test_logging.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypedDict
 
 import pytest
 
@@ -11,57 +11,65 @@
 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
 
 
-class TestLoggingStore(StoreTests[LoggingStore, cpu.Buffer]):
-    store_cls = LoggingStore
+class StoreKwargs(TypedDict):
+    store: LocalStore
+    log_level: str
+
+
+class TestLoggingStore(StoreTests[LoggingStore[LocalStore], cpu.Buffer]):
+    # store_cls is needed to do an isintsance check, so can't be a subscripted generic
+    store_cls = LoggingStore  # type: ignore[assignment]
     buffer_cls = cpu.Buffer
 
-    async def get(self, store: LoggingStore, key: str) -> Buffer:
+    async def get(self, store: LoggingStore[LocalStore], key: str) -> Buffer:
         return self.buffer_cls.from_bytes((store._store.root / key).read_bytes())
 
-    async def set(self, store: LoggingStore, key: str, value: Buffer) -> None:
+    async def set(self, store: LoggingStore[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)), "log_level": "DEBUG"}
+    def store_kwargs(self, tmp_path: Path) -> StoreKwargs:
+        return {"store": LocalStore(str(tmp_path)), "log_level": "DEBUG"}
 
     @pytest.fixture
-    def open_kwargs(self, tmpdir) -> dict[str, str]:
-        return {"store_cls": LocalStore, "root": str(tmpdir), "log_level": "DEBUG"}
+    def open_kwargs(self, tmp_path: Path) -> dict[str, type[LocalStore] | str]:
+        return {"store_cls": LocalStore, "root": str(tmp_path), "log_level": "DEBUG"}
 
     @pytest.fixture
-    def store(self, store_kwargs: str | dict[str, Buffer] | None) -> LoggingStore:
+    def store(self, store_kwargs: StoreKwargs) -> LoggingStore[LocalStore]:
         return self.store_cls(**store_kwargs)
 
-    def test_store_supports_writes(self, store: LoggingStore) -> None:
+    def test_store_supports_writes(self, store: LoggingStore[LocalStore]) -> None:
         assert store.supports_writes
 
-    def test_store_supports_listing(self, store: LoggingStore) -> None:
+    def test_store_supports_listing(self, store: LoggingStore[LocalStore]) -> None:
         assert store.supports_listing
 
-    def test_store_repr(self, store: LoggingStore) -> None:
+    def test_store_repr(self, store: LoggingStore[LocalStore]) -> None:
         assert f"{store!r}" == f"LoggingStore(LocalStore, 'file://{store._store.root.as_posix()}')"
 
-    def test_store_str(self, store: LoggingStore) -> None:
+    def test_store_str(self, store: LoggingStore[LocalStore]) -> None:
         assert str(store) == f"logging-file://{store._store.root.as_posix()}"
 
-    async def test_default_handler(self, local_store, capsys) -> None:
+    async def test_default_handler(
+        self, local_store: LocalStore, capsys: pytest.CaptureFixture[str]
+    ) -> None:
         # Store and then remove existing handlers to enter default handler code path
         handlers = logging.getLogger().handlers[:]
         for h in handlers:
             logging.getLogger().removeHandler(h)
         # Test logs are sent to stdout
         wrapped = LoggingStore(store=local_store)
         buffer = default_buffer_prototype().buffer
-        res = await wrapped.set("foo/bar/c/0", buffer.from_bytes(b"\x01\x02\x03\x04"))
+        res = await wrapped.set("foo/bar/c/0", buffer.from_bytes(b"\x01\x02\x03\x04"))  # type: ignore[func-returns-value]
         assert res is None
         captured = capsys.readouterr()
         assert len(captured) == 2
@@ -71,7 +79,7 @@ async def test_default_handler(self, local_store, capsys) -> None:
         for h in handlers:
             logging.getLogger().addHandler(h)
 
-    def test_is_open_setter_raises(self, store: LoggingStore) -> None:
+    def test_is_open_setter_raises(self, store: LoggingStore[LocalStore]) -> None:
         "Test that a user cannot change `_is_open` without opening the underlying store."
         with pytest.raises(
             NotImplementedError, match="LoggingStore must be opened via the `_open` method"
@@ -80,12 +88,12 @@ def test_is_open_setter_raises(self, store: LoggingStore) -> None:
 
 
 @pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=["store"])
-async def test_logging_store(store: Store, caplog) -> None:
+async def test_logging_store(store: Store, caplog: pytest.LogCaptureFixture) -> None:
     wrapped = LoggingStore(store=store, log_level="DEBUG")
     buffer = default_buffer_prototype().buffer
 
     caplog.clear()
-    res = await wrapped.set("foo/bar/c/0", buffer.from_bytes(b"\x01\x02\x03\x04"))
+    res = await wrapped.set("foo/bar/c/0", buffer.from_bytes(b"\x01\x02\x03\x04"))  # type: ignore[func-returns-value]
     assert res is None
     assert len(caplog.record_tuples) == 2
     for tup in caplog.record_tuples: