8000 Fix typing in a bunch of store tests (#3052) · zarr-developers/zarr-python@7584b96 · GitHub
[go: up one dir, main page]
Skip to content

Commit 7584b96

Browse files
dstansbyd-v-b
dstansby
and
d-v-b
authored
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>
1 parent 520bc1f commit 7584b96

File tree

8 files changed

+80
-52
lines changed

8 files changed

+80
-52
lines changed

‎.pre-commit-config.yaml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ repos:
3838
# Tests
3939
- pytest
4040
- hypothesis
41+
- s3fs
4142
- repo: https://github.com/scientific-python/cookie
4243
rev: 2025.01.22
4344
hooks:

‎pyproject.toml

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -358,7 +358,11 @@ module = [
358358
"tests.package_with_entrypoint.*",
359359
"zarr.testing.stateful",
360360
"tests.test_codecs.test_transpose",
361-
"tests.test_config"
361+
"tests.test_config",
362+
"tests.test_store.test_zip",
363+
"tests.test_store.test_local",
364+
"tests.test_store.test_fsspec",
365+
"tests.test_store.test_memory",
362366
]
363367
strict = false
364368

@@ -368,7 +372,11 @@ strict = false
368372
module = [
369373
"tests.test_codecs.test_codecs",
370374
"tests.test_metadata.*",
371-
"tests.test_store.*",
375+
"tests.test_store.test_core",
376+
"tests.test_store.test_logging",
377+
"tests.test_store.test_object",
378+
"tests.test_store.test_stateful",
379+
"tests.test_store.test_wrapper",
372380
"tests.test_group",
373381
"tests.test_indexing",
374382
"tests.test_properties",

‎src/zarr/testing/store.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ async def get(self, store: S, key: str) -> Buffer:
5858

5959
@abstractmethod
6060
@pytest.fixture
61-
def store_kwargs(self) -> dict[str, Any]:
61+
def store_kwargs(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
6262
"""Kwargs for instantiating a store"""
6363
...
6464

‎src/zarr/testing/utils.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
from __future__ import annotations
22

3-
from collections.abc import Callable, Coroutine
4-
from typing import TYPE_CHECKING, Any, TypeVar, cast
3+
from typing import TYPE_CHECKING, TypeVar, cast
54

65
import pytest
76

@@ -38,13 +37,13 @@ def has_cupy() -> bool:
3837
return False
3938

4039

41-
T_Callable = TypeVar("T_Callable", bound=Callable[..., Coroutine[Any, Any, None] | None])
40+
T = TypeVar("T")
4241

4342

4443
# Decorator for GPU tests
45-
def gpu_test(func: T_Callable) -> T_Callable:
44+
def gpu_test(func: T) -> T:
4645
return cast(
47-
"T_Callable",
46+
"T",
4847
pytest.mark.gpu(
4948
pytest.mark.skipif(not has_cupy(), reason="CuPy not installed or no GPU available")(
5049
func

‎tests/test_store/test_fsspec.py

Lines changed: 25 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import json
44
import os
55
import re
6-
from typing import TYPE_CHECKING
6+
from typing import TYPE_CHECKING, Any
77

88
import pytest
99
from packaging.version import parse as parse_version
@@ -17,8 +17,13 @@
1717

1818
if TYPE_CHECKING:
1919
from collections.abc import Generator
20+
from pathlib import Path
2021

2122
import botocore.client
23+
import s3fs
24+
25+
from zarr.core.common import JSON
26+
2227

2328
# Warning filter due to https://github.com/boto/boto3/issues/3889
2429
pytestmark = [
@@ -109,10 +114,13 @@ async def test_basic() -> None:
109114
data = b"hello"
110115
await store.set("foo", cpu.Buffer.from_bytes(data))
111116
assert await store.exists("foo")
112-
assert (await store.get("foo", prototype=default_buffer_prototype())).to_bytes() == data
117+
buf = await store.get("foo", prototype=default_buffer_prototype())
118+
assert buf is not None
119+
assert buf.to_bytes() == data
113120
out = await store.get_partial_values(
114121
prototype=default_buffer_prototype(), key_ranges=[("foo", OffsetByteRequest(1))]
115122
)
123+
assert out[0] is not None
116124
assert out[0].to_bytes() == data[1:]
117125

118126

@@ -121,7 +129,7 @@ class TestFsspecStoreS3(StoreTests[FsspecStore, cpu.Buffer]):
121129
buffer_cls = cpu.Buffer
122130

123131
@pytest.fixture
124-
def store_kwargs(self, request) -> dict[str, str | bool]:
132+
def store_kwargs(self) -> dict[str, str | bool]:
125133
try:
126134
from fsspec import url_to_fs
127135
except ImportError:
@@ -133,7 +141,7 @@ def store_kwargs(self, request) -> dict[str, str | bool]:
133141
return {"fs": fs, "path": path}
134142

135143
@pytest.fixture
136-
def store(self, store_kwargs: dict[str, str | bool]) -> FsspecStore:
144+
async def store(self, store_kwargs: dict[str, Any]) -> FsspecStore:
137145
return self.store_cls(**store_kwargs)
138146

139147
async def get(self, store: FsspecStore, key: str) -> Buffer:
@@ -168,7 +176,11 @@ async def test_fsspec_store_from_uri(self, store: FsspecStore) -> None:
168176
"anon": False,
169177
}
170178

171-
meta = {"attributes": {"key": "value"}, "zarr_format": 3, "node_type": "group"}
179+
meta: dict[str, JSON] = {
180+
"attributes": {"key": "value"},
181+
"zarr_format": 3,
182+
"node_type": "group",
183+
}
172184

173185
await store.set(
174186
"zarr.json",
@@ -179,7 +191,7 @@ async def test_fsspec_store_from_uri(self, store: FsspecStore) -> None:
179191
)
180192
assert dict(group.attrs) == {"key": "value"}
181193

182-
meta["attributes"]["key"] = "value-2"
194+
meta["attributes"]["key"] = "value-2" # type: ignore[index]
183195
await store.set(
184196
"directory-2/zarr.json",
185197
self.buffer_cls.from_bytes(json.dumps(meta).encode()),
@@ -189,7 +201,7 @@ async def test_fsspec_store_from_uri(self, store: FsspecStore) -> None:
189201
)
190202
assert dict(group.attrs) == {"key": "value-2"}
191203

192-
meta["attributes"]["key"] = "value-3"
204+
meta["attributes"]["key"] = "value-3" # type: ignore[index]
193205
await store.set(
194206
"directory-3/zarr.json",
195207
self.buffer_cls.from_bytes(json.dumps(meta).encode()),
@@ -216,7 +228,7 @@ def test_from_upath(self) -> None:
216228
assert result.fs.asynchronous
217229
assert result.path == f"{test_bucket_name}/foo/bar"
218230

219-
def test_init_raises_if_path_has_scheme(self, store_kwargs) -> None:
231+
def test_init_raises_if_path_has_scheme(self, store_kwargs: dict[str, Any]) -> None:
220232
# regression test for https://github.com/zarr-developers/zarr-python/issues/2342
221233
store_kwargs["path"] = "s3://" + store_kwargs["path"]
222234
with pytest.raises(
@@ -237,7 +249,7 @@ def test_init_warns_if_fs_asynchronous_is_false(self) -> None:
237249
with pytest.warns(UserWarning, match=r".* was not created with `asynchronous=True`.*"):
238250
self.store_cls(**store_kwargs)
239251

240-
async def test_empty_nonexistent_path(self, store_kwargs) -> None:
252+
async def test_empty_nonexistent_path(self, store_kwargs: dict[str, Any]) -> None:
241253
# regression test for https://github.com/zarr-developers/zarr-python/pull/2343
242254
store_kwargs["path"] += "/abc"
243255
store = await self.store_cls.open(**store_kwargs)
@@ -256,7 +268,7 @@ async def test_delete_dir_unsupported_deletes(self, store: FsspecStore) -> None:
256268
parse_version(fsspec.__version__) < parse_version("2024.12.0"),
257269
reason="No AsyncFileSystemWrapper",
258270
)
259-
def test_wrap_sync_filesystem():
271+
def test_wrap_sync_filesystem() -> None:
260272
"""The local fs is not async so we should expect it to be wrapped automatically"""
261273
from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
262274

@@ -270,7 +282,7 @@ def test_wrap_sync_filesystem():
270282
parse_version(fsspec.__version__) < parse_version("2024.12.0"),
271283
reason="No AsyncFileSystemWrapper",
272284
)
273-
def test_no_wrap_async_filesystem():
285+
def test_no_wrap_async_filesystem() -> None:
274286
"""An async fs should not be wrapped automatically; fsspec's https filesystem is such an fs"""
275287
from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
276288

@@ -284,12 +296,12 @@ def test_no_wrap_async_filesystem():
284296
parse_version(fsspec.__version__) < parse_version("2024.12.0"),
285297
reason="No AsyncFileSystemWrapper",
286298
)
287-
async def test_delete_dir_wrapped_filesystem(tmpdir) -> None:
299+
async def test_delete_dir_wrapped_filesystem(tmp_path: Path) -> None:
288300
from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
289301
from fsspec.implementations.local import LocalFileSystem
290302

291303
wrapped_fs = AsyncFileSystemWrapper(LocalFileSystem(auto_mkdir=True))
292-
store = FsspecStore(wrapped_fs, read_only=False, path=f"{tmpdir}/test/path")
304+
store = FsspecStore(wrapped_fs, read_only=False, path=f"{tmp_path}/test/path")
293305

294306
assert isinstance(store.fs, AsyncFileSystemWrapper)
295307
assert store.fs.asynchronous

‎tests/test_store/test_local.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ async def set(self, store: LocalStore, key: str, value: Buffer) -> None:
2828
(store.root / key).write_bytes(value.to_bytes())
2929

3030
@pytest.fixture
31-
def store_kwargs(self, tmpdir) -> dict[str, str]:
31+
def store_kwargs(self, tmpdir: str) -> dict[str, str]:
3232
return {"root": str(tmpdir)}
3333

3434
def test_store_repr(self, store: LocalStore) -> None:
@@ -48,24 +48,24 @@ async def test_empty_with_empty_subdir(self, store: LocalStore) -> None:
4848
(store.root / "foo/bar").mkdir(parents=True)
4949
assert await store.is_empty("")
5050

51-
def test_creates_new_directory(self, tmp_path: pathlib.Path):
51+
def test_creates_new_directory(self, tmp_path: pathlib.Path) -> None:
5252
target = tmp_path.joinpath("a", "b", "c")
5353
assert not target.exists()
5454

5555
store = self.store_cls(root=target)
5656
zarr.group(store=store)
5757

58-
def test_invalid_root_raises(self):
58+
def test_invalid_root_raises(self) -> None:
5959
"""
6060
Test that a TypeError is raised when a non-str/Path type is used for the `root` argument
6161
"""
6262
with pytest.raises(
6363
TypeError,
6464
match=r"'root' must be a string or Path instance. Got an instance of <class 'int'> instead.",
6565
):
66-
LocalStore(root=0)
66+
LocalStore(root=0) # type: ignore[arg-type]
6767

68-
async def test_get_with_prototype_default(self, store: LocalStore):
68+
async def test_get_with_prototype_default(self, store: LocalStore) -> None:
6969
"""
7070
Ensure that data can be read via ``store.get`` if the prototype keyword argument is unspecified, i.e. set to ``None``.
7171
"""

‎tests/test_store/test_memory.py

Lines changed: 22 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
11
from __future__ import annotations
22

33
import re
4-
from typing import TYPE_CHECKING
4+
from typing import TYPE_CHECKING, Any
55

66
import numpy as np
7+
import numpy.typing as npt
78
import pytest
89

910
import zarr
11+
import zarr.core
12+
import zarr.core.array
1013
from zarr.core.buffer import Buffer, cpu, gpu
1114
from zarr.storage import GpuMemoryStore, MemoryStore
1215
from zarr.testing.store import StoreTests
@@ -31,16 +34,16 @@ async def get(self, store: MemoryStore, key: str) -> Buffer:
3134
return store._store_dict[key]
3235

3336
@pytest.fixture(params=[None, True])
34-
def store_kwargs(
35-
self, request: pytest.FixtureRequest
36-
) -> dict[str, str | dict[str, Buffer] | None]:
37-
kwargs = {"store_dict": None}
37+
def store_kwargs(self, request: pytest.FixtureRequest) -> dict[str, Any]:
38+
kwargs: dict[str, Any]
3839
if request.param is True:
39-
kwargs["store_dict"] = {}
40+
kwargs = {"store_dict": {}}
41+
else:
42+
kwargs = {"store_dict": None}
4043
return kwargs
4144

4245
@pytest.fixture
43-
def store(self, store_kwargs: str | dict[str, Buffer] | None) -> MemoryStore:
46+
async def store(self, store_kwargs: dict[str, Any]) -> MemoryStore:
4447
return self.store_cls(**store_kwargs)
4548

4649
def test_store_repr(self, store: MemoryStore) -> None:
@@ -55,13 +58,13 @@ def test_store_supports_listing(self, store: MemoryStore) -> None:
5558
def test_store_supports_partial_writes(self, store: MemoryStore) -> None:
5659
assert store.supports_partial_writes
5760

58-
def test_list_prefix(self, store: MemoryStore) -> None:
61+
async def test_list_prefix(self, store: MemoryStore) -> None:
5962
assert True
6063

6164
@pytest.mark.parametrize("dtype", ["uint8", "float32", "int64"])
6265
@pytest.mark.parametrize("zarr_format", [2, 3])
6366
async def test_deterministic_size(
64-
self, store: MemoryStore, dtype, zarr_format: ZarrFormat
67+
self, store: MemoryStore, dtype: npt.DTypeLike, zarr_format: ZarrFormat
6568
) -> None:
6669
a = zarr.empty(
6770
store=store,
@@ -85,23 +88,23 @@ class TestGpuMemoryStore(StoreTests[GpuMemoryStore, gpu.Buffer]):
8588
store_cls = GpuMemoryStore
8689
buffer_cls = gpu.Buffer
8790

88-
async def set(self, store: GpuMemoryStore, key: str, value: Buffer) -> None:
91+
async def set(self, store: GpuMemoryStore, key: str, value: gpu.Buffer) -> None: # type: ignore[override]
8992
store._store_dict[key] = value
9093

9194
async def get(self, store: MemoryStore, key: str) -> Buffer:
9295
return store._store_dict[key]
9396

9497
@pytest.fixture(params=[None, True])
95-
def store_kwargs(
96-
self, request: pytest.FixtureRequest
97-
) -> dict[str, str | dict[str, Buffer] | None]:
98-
kwargs = {"store_dict": None}
98+
def store_kwargs(self, request: pytest.FixtureRequest) -> dict[str, Any]:
99+
kwargs: dict[str, Any]
99100
if request.param is True:
100-
kwargs["store_dict"] = {}
101+
kwargs = {"store_dict": {}}
102+
else:
103+
kwargs = {"store_dict": None}
101104
return kwargs
102105

103106
@pytest.fixture
104-
def store(self, store_kwargs: str | dict[str, gpu.Buffer] | None) -> GpuMemoryStore:
107+
async def store(self, store_kwargs: dict[str, Any]) -> GpuMemoryStore:
105108
return self.store_cls(**store_kwargs)
106109

107110
def test_store_repr(self, store: GpuMemoryStore) -> None:
@@ -116,15 +119,15 @@ def test_store_supports_listing(self, store: GpuMemoryStore) -> None:
116119
def test_store_supports_partial_writes(self, store: GpuMemoryStore) -> None:
117120
assert store.supports_partial_writes
118121

119-
def test_list_prefix(self, store: GpuMemoryStore) -> None:
122+
async def test_list_prefix(self, store: GpuMemoryStore) -> None:
120123
assert True
121124

122125
def test_dict_reference(self, store: GpuMemoryStore) -> None:
123-
store_dict = {}
126+
store_dict: dict[str, Any] = {}
124127
result = GpuMemoryStore(store_dict=store_dict)
125128
assert result._store_dict is store_dict
126129

127-
def test_from_dict(self):
130+
def test_from_dict(self) -> None:
128131
d = {
129132
"a": gpu.Buffer.from_bytes(b"aaaa"),
130133
"b": cpu.Buffer.from_bytes(b"bbbb"),

0 commit comments

Comments
 (0)
0