Fix typing in a bunch of store tests (#3052)

* Ignore explicit test store files

* Fix test_zip

* Fix test_local

* Fix test_fsspec

* Fix typing in test_memory

* Remove walrus

Co-authored-by: Davis Bennett <davis.v.bennett@gmail.com>

---------

Co-authored-by: Davis Bennett <davis.v.bennett@gmail.com>

Author: dstansby

.pre-commit-config.yaml

@@ -38,6 +38,7 @@ repos:
           # Tests
           - pytest
           - hypothesis
+          - s3fs
   - repo: https://github.com/scientific-python/cookie
     rev: 2025.01.22
     hooks:

pyproject.toml

@@ -358,7 +358,11 @@ module = [
     "tests.package_with_entrypoint.*",
     "zarr.testing.stateful",
     "tests.test_codecs.test_transpose",
-    "tests.test_config"
+    "tests.test_config",
+    "tests.test_store.test_zip",
+    "tests.test_store.test_local",
+    "tests.test_store.test_fsspec",
+    "tests.test_store.test_memory",
 ]
 strict = false
 
@@ -368,7 +372,11 @@ strict = false
 module = [
     "tests.test_codecs.test_codecs",
     "tests.test_metadata.*",
-    "tests.test_store.*",
+    "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/testing/store.py

@@ -58,7 +58,7 @@ async def get(self, store: S, key: str) -> Buffer:
 
     @abstractmethod
     @pytest.fixture
-    def store_kwargs(self) -> dict[str, Any]:
+    def store_kwargs(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
         """Kwargs for instantiating a store"""
         ...
 

src/zarr/testing/utils.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Callable, Coroutine
-from typing import TYPE_CHECKING, Any, TypeVar, cast
+from typing import TYPE_CHECKING, TypeVar, cast
 
 import pytest
 
@@ -38,13 +37,13 @@ def has_cupy() -> bool:
         return False
 
 
-T_Callable = TypeVar("T_Callable", bound=Callable[..., Coroutine[Any, Any, None] | None])
+T = TypeVar("T")
 
 
 # Decorator for GPU tests
-def gpu_test(func: T_Callable) -> T_Callable:
+def gpu_test(func: T) -> T:
     return cast(
-        "T_Callable",
+        "T",
         pytest.mark.gpu(
             pytest.mark.skipif(not has_cupy(), reason="CuPy not installed or no GPU available")(
                 func

tests/test_store/test_fsspec.py

@@ -3,7 +3,7 @@
 import json
 import os
 import re
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import pytest
 from packaging.version import parse as parse_version
@@ -17,8 +17,13 @@
 
 if TYPE_CHECKING:
     from collections.abc import Generator
+    from pathlib import Path
 
     import botocore.client
+    import s3fs
+
+    from zarr.core.common import JSON
+
 
 # Warning filter due to https://github.com/boto/boto3/issues/3889
 pytestmark = [
@@ -109,10 +114,13 @@ async def test_basic() -> None:
     data = b"hello"
     await store.set("foo", cpu.Buffer.from_bytes(data))
     assert await store.exists("foo")
-    assert (await store.get("foo", prototype=default_buffer_prototype())).to_bytes() == data
+    buf = await store.get("foo", prototype=default_buffer_prototype())
+    assert buf is not None
+    assert buf.to_bytes() == data
     out = await store.get_partial_values(
         prototype=default_buffer_prototype(), key_ranges=[("foo", OffsetByteRequest(1))]
     )
+    assert out[0] is not None
     assert out[0].to_bytes() == data[1:]
 
 
@@ -121,7 +129,7 @@ class TestFsspecStoreS3(StoreTests[FsspecStore, cpu.Buffer]):
     buffer_cls = cpu.Buffer
 
     @pytest.fixture
-    def store_kwargs(self, request) -> dict[str, str | bool]:
+    def store_kwargs(self) -> dict[str, str | bool]:
         try:
             from fsspec import url_to_fs
         except ImportError:
@@ -133,7 +141,7 @@ def store_kwargs(self, request) -> dict[str, str | bool]:
         return {"fs": fs, "path": path}
 
     @pytest.fixture
-    def store(self, store_kwargs: dict[str, str | bool]) -> FsspecStore:
+    async def store(self, store_kwargs: dict[str, Any]) -> FsspecStore:
         return self.store_cls(**store_kwargs)
 
     async def get(self, store: FsspecStore, key: str) -> Buffer:
@@ -168,7 +176,11 @@ async def test_fsspec_store_from_uri(self, store: FsspecStore) -> None:
             "anon": False,
         }
 
-        meta = {"attributes": {"key": "value"}, "zarr_format": 3, "node_type": "group"}
+        meta: dict[str, JSON] = {
+            "attributes": {"key": "value"},
+            "zarr_format": 3,
+            "node_type": "group",
+        }
 
         await store.set(
             "zarr.json",
@@ -179,7 +191,7 @@ async def test_fsspec_store_from_uri(self, store: FsspecStore) -> None:
         )
         assert dict(group.attrs) == {"key": "value"}
 
-        meta["attributes"]["key"] = "value-2"
+        meta["attributes"]["key"] = "value-2"  # type: ignore[index]
         await store.set(
             "directory-2/zarr.json",
             self.buffer_cls.from_bytes(json.dumps(meta).encode()),
@@ -189,7 +201,7 @@ async def test_fsspec_store_from_uri(self, store: FsspecStore) -> None:
         )
         assert dict(group.attrs) == {"key": "value-2"}
 
-        meta["attributes"]["key"] = "value-3"
+        meta["attributes"]["key"] = "value-3"  # type: ignore[index]
         await store.set(
             "directory-3/zarr.json",
             self.buffer_cls.from_bytes(json.dumps(meta).encode()),
@@ -216,7 +228,7 @@ def test_from_upath(self) -> None:
         assert result.fs.asynchronous
         assert result.path == f"{test_bucket_name}/foo/bar"
 
-    def test_init_raises_if_path_has_scheme(self, store_kwargs) -> None:
+    def test_init_raises_if_path_has_scheme(self, store_kwargs: dict[str, Any]) -> None:
         # regression test for https://github.com/zarr-developers/zarr-python/issues/2342
         store_kwargs["path"] = "s3://" + store_kwargs["path"]
         with pytest.raises(
@@ -237,7 +249,7 @@ def test_init_warns_if_fs_asynchronous_is_false(self) -> None:
         with pytest.warns(UserWarning, match=r".* was not created with `asynchronous=True`.*"):
             self.store_cls(**store_kwargs)
 
-    async def test_empty_nonexistent_path(self, store_kwargs) -> None:
+    async def test_empty_nonexistent_path(self, store_kwargs: dict[str, Any]) -> None:
         # regression test for https://github.com/zarr-developers/zarr-python/pull/2343
         store_kwargs["path"] += "/abc"
         store = await self.store_cls.open(**store_kwargs)
@@ -256,7 +268,7 @@ async def test_delete_dir_unsupported_deletes(self, store: FsspecStore) -> None:
     parse_version(fsspec.__version__) < parse_version("2024.12.0"),
     reason="No AsyncFileSystemWrapper",
 )
-def test_wrap_sync_filesystem():
+def test_wrap_sync_filesystem() -> None:
     """The local fs is not async so we should expect it to be wrapped automatically"""
     from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
 
@@ -270,7 +282,7 @@ def test_wrap_sync_filesystem():
     parse_version(fsspec.__version__) < parse_version("2024.12.0"),
     reason="No AsyncFileSystemWrapper",
 )
-def test_no_wrap_async_filesystem():
+def test_no_wrap_async_filesystem() -> None:
     """An async fs should not be wrapped automatically; fsspec's https filesystem is such an fs"""
     from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
 
@@ -284,12 +296,12 @@ def test_no_wrap_async_filesystem():
     parse_version(fsspec.__version__) < parse_version("2024.12.0"),
     reason="No AsyncFileSystemWrapper",
 )
-async def test_delete_dir_wrapped_filesystem(tmpdir) -> None:
+async def test_delete_dir_wrapped_filesystem(tmp_path: Path) -> None:
     from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
     from fsspec.implementations.local import LocalFileSystem
 
     wrapped_fs = AsyncFileSystemWrapper(LocalFileSystem(auto_mkdir=True))
-    store = FsspecStore(wrapped_fs, read_only=False, path=f"{tmpdir}/test/path")
+    store = FsspecStore(wrapped_fs, read_only=False, path=f"{tmp_path}/test/path")
 
     assert isinstance(store.fs, AsyncFileSystemWrapper)
     assert store.fs.asynchronous

tests/test_store/test_local.py

@@ -28,7 +28,7 @@ async def set(self, store: LocalStore, key: str, value: Buffer) -> None:
         (store.root / key).write_bytes(value.to_bytes())
 
     @pytest.fixture
-    def store_kwargs(self, tmpdir) -> dict[str, str]:
+    def store_kwargs(self, tmpdir: str) -> dict[str, str]:
         return {"root": str(tmpdir)}
 
     def test_store_repr(self, store: LocalStore) -> None:
@@ -48,24 +48,24 @@ async def test_empty_with_empty_subdir(self, store: LocalStore) -> None:
         (store.root / "foo/bar").mkdir(parents=True)
         assert await store.is_empty("")
 
-    def test_creates_new_directory(self, tmp_path: pathlib.Path):
+    def test_creates_new_directory(self, tmp_path: pathlib.Path) -> None:
         target = tmp_path.joinpath("a", "b", "c")
         assert not target.exists()
 
         store = self.store_cls(root=target)
         zarr.group(store=store)
 
-    def test_invalid_root_raises(self):
+    def test_invalid_root_raises(self) -> None:
         """
         Test that a TypeError is raised when a non-str/Path type is used for the `root` argument
         """
         with pytest.raises(
             TypeError,
             match=r"'root' must be a string or Path instance. Got an instance of <class 'int'> instead.",
         ):
-            LocalStore(root=0)
+            LocalStore(root=0)  # type: ignore[arg-type]
 
-    async def test_get_with_prototype_default(self, store: LocalStore):
+    async def test_get_with_prototype_default(self, store: LocalStore) -> None:
         """
         Ensure that data can be read via ``store.get`` if the prototype keyword argument is unspecified, i.e. set to ``None``.
         """

tests/test_store/test_memory.py

@@ -1,12 +1,15 @@
 from __future__ import annotations
 
 import re
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
+import numpy.typing as npt
 import pytest
 
 import zarr
+import zarr.core
+import zarr.core.array
 from zarr.core.buffer import Buffer, cpu, gpu
 from zarr.storage import GpuMemoryStore, MemoryStore
 from zarr.testing.store import StoreTests
@@ -31,16 +34,16 @@ async def get(self, store: MemoryStore, key: str) -> Buffer:
         return store._store_dict[key]
 
     @pytest.fixture(params=[None, True])
-    def store_kwargs(
-        self, request: pytest.FixtureRequest
-    ) -> dict[str, str | dict[str, Buffer] | None]:
-        kwargs = {"store_dict": None}
+    def store_kwargs(self, request: pytest.FixtureRequest) -> dict[str, Any]:
+        kwargs: dict[str, Any]
         if request.param is True:
-            kwargs["store_dict"] = {}
+            kwargs = {"store_dict": {}}
+        else:
+            kwargs = {"store_dict": None}
         return kwargs
 
     @pytest.fixture
-    def store(self, store_kwargs: str | dict[str, Buffer] | None) -> MemoryStore:
+    async def store(self, store_kwargs: dict[str, Any]) -> MemoryStore:
         return self.store_cls(**store_kwargs)
 
     def test_store_repr(self, store: MemoryStore) -> None:
@@ -55,13 +58,13 @@ def test_store_supports_listing(self, store: MemoryStore) -> None:
     def test_store_supports_partial_writes(self, store: MemoryStore) -> None:
         assert store.supports_partial_writes
 
-    def test_list_prefix(self, store: MemoryStore) -> None:
+    async def test_list_prefix(self, store: MemoryStore) -> None:
         assert True
 
     @pytest.mark.parametrize("dtype", ["uint8", "float32", "int64"])
     @pytest.mark.parametrize("zarr_format", [2, 3])
     async def test_deterministic_size(
-        self, store: MemoryStore, dtype, zarr_format: ZarrFormat
+        self, store: MemoryStore, dtype: npt.DTypeLike, zarr_format: ZarrFormat
     ) -> None:
         a = zarr.empty(
             store=store,
@@ -85,23 +88,23 @@ class TestGpuMemoryStore(StoreTests[GpuMemoryStore, gpu.Buffer]):
     store_cls = GpuMemoryStore
     buffer_cls = gpu.Buffer
 
-    async def set(self, store: GpuMemoryStore, key: str, value: Buffer) -> None:
+    async def set(self, store: GpuMemoryStore, key: str, value: gpu.Buffer) -> None:  # type: ignore[override]
         store._store_dict[key] = value
 
     async def get(self, store: MemoryStore, key: str) -> Buffer:
         return store._store_dict[key]
 
     @pytest.fixture(params=[None, True])
-    def store_kwargs(
-        self, request: pytest.FixtureRequest
-    ) -> dict[str, str | dict[str, Buffer] | None]:
-        kwargs = {"store_dict": None}
+    def store_kwargs(self, request: pytest.FixtureRequest) -> dict[str, Any]:
+        kwargs: dict[str, Any]
         if request.param is True:
-            kwargs["store_dict"] = {}
+            kwargs = {"store_dict": {}}
+        else:
+            kwargs = {"store_dict": None}
         return kwargs
 
     @pytest.fixture
-    def store(self, store_kwargs: str | dict[str, gpu.Buffer] | None) -> GpuMemoryStore:
+    async def store(self, store_kwargs: dict[str, Any]) -> GpuMemoryStore:
         return self.store_cls(**store_kwargs)
 
     def test_store_repr(self, store: GpuMemoryStore) -> None:
@@ -116,15 +119,15 @@ def test_store_supports_listing(self, store: GpuMemoryStore) -> None:
     def test_store_supports_partial_writes(self, store: GpuMemoryStore) -> None:
         assert store.supports_partial_writes
 
-    def test_list_prefix(self, store: GpuMemoryStore) -> None:
+    async def test_list_prefix(self, store: GpuMemoryStore) -> None:
         assert True
 
     def test_dict_reference(self, store: GpuMemoryStore) -> None:
-        store_dict = {}
+        store_dict: dict[str, Any] = {}
         result = GpuMemoryStore(store_dict=store_dict)
         assert result._store_dict is store_dict
 
-    def test_from_dict(self):
+    def test_from_dict(self) -> None:
         d = {
             "a": gpu.Buffer.from_bytes(b"aaaa"),
             "b": cpu.Buffer.from_bytes(b"bbbb"),