Fix typing in test_memory

Author: dstansby

pyproject.toml

@@ -361,7 +361,8 @@ module = [
     "tests.test_config",
     "tests.test_store.test_zip",
     "tests.test_store.test_local",
-    "tests.test_store.test_fsspec"
+    "tests.test_store.test_fsspec",
+    "tests.test_store.test_memory",
 ]
 strict = false
 
@@ -373,7 +374,6 @@ module = [
     "tests.test_metadata.*",
     "tests.test_store.test_core",
     "tests.test_store.test_logging",
-    "tests.test_store.test_memory",
     "tests.test_store.test_object",
     "tests.test_store.test_stateful",
     "tests.test_store.test_wrapper",

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_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"),