Check tests folder with mypy (#2150)

* Check tests folder with mypy

* use shapelike

* fixup

Author: TomAugspurger

.pre-commit-config.yaml

@@ -25,7 +25,7 @@ repos:
     rev: v1.11.2
     hooks:
       - id: mypy
-        files: src
+        files: src|tests/v3/test_(api|array|buffer).py
         additional_dependencies:
           # Package dependencies
           - asciitree

src/zarr/core/array.py

@@ -25,8 +25,10 @@
     ZARRAY_JSON,
     ZATTRS_JSON,
     ChunkCoords,
+    ShapeLike,
     ZarrFormat,
     concurrent_map,
+    parse_shapelike,
     product,
 )
 from zarr.core.config import config, parse_indexing_order
@@ -116,7 +118,7 @@ async def create(
         store: StoreLike,
         *,
         # v2 and v3
-        shape: ChunkCoords,
+        shape: ShapeLike,
         dtype: npt.DTypeLike,
         zarr_format: ZarrFormat = 3,
         fill_value: Any | None = None,
@@ -132,7 +134,7 @@ async def create(
         codecs: Iterable[Codec | dict[str, JSON]] | None = None,
         dimension_names: Iterable[str] | None = None,
         # v2 only
-        chunks: ChunkCoords | None = None,
+        chunks: ShapeLike | None = None,
         dimension_separator: Literal[".", "/"] | None = None,
         order: Literal["C", "F"] | None = None,
         filters: list[dict[str, JSON]] | None = None,
@@ -143,9 +145,14 @@ async def create(
     ) -> AsyncArray:
         store_path = await make_store_path(store)
 
+        shape = parse_shapelike(shape)
+
         if chunk_shape is None:
             if chunks is None:
                 chunk_shape = chunks = _guess_chunks(shape=shape, typesize=np.dtype(dtype).itemsize)
+            else:
+                chunks = parse_shapelike(chunks)
+
             chunk_shape = chunks
         elif chunks is not None:
             raise ValueError("Only one of chunk_shape or chunks must be provided.")
@@ -217,7 +224,7 @@ async def _create_v3(
         cls,
         store_path: StorePath,
         *,
-        shape: ChunkCoords,
+        shape: ShapeLike,
         dtype: npt.DTypeLike,
         chunk_shape: ChunkCoords,
         fill_value: Any | None = None,
@@ -235,6 +242,7 @@ async def _create_v3(
         if not exists_ok:
             await ensure_no_existing_node(store_path, zarr_format=3)
 
+        shape = parse_shapelike(shape)
         codecs = list(codecs) if codecs is not None else [BytesCodec()]
 
         if fill_value is None:

src/zarr/core/chunk_grids.py

@@ -15,6 +15,7 @@
     JSON,
     ChunkCoords,
     ChunkCoordsLike,
+    ShapeLike,
     parse_named_configuration,
     parse_shapelike,
 )
@@ -27,7 +28,7 @@
 
 
 def _guess_chunks(
-    shape: ChunkCoords,
+    shape: ShapeLike,
     typesize: int,
     *,
     increment_bytes: int = 256 * 1024,
@@ -57,6 +58,8 @@ def _guess_chunks(
     ChunkCoords
 
     """
+    if isinstance(shape, int):
+        shape = (shape,)
 
     ndims = len(shape)
     # require chunks to have non-zero length for all dimensions

src/zarr/core/group.py

@@ -22,6 +22,7 @@
     ZATTRS_JSON,
     ZGROUP_JSON,
     ChunkCoords,
+    ShapeLike,
     ZarrFormat,
     parse_shapelike,
 )
@@ -365,7 +366,7 @@ async def create_array(
         self,
         name: str,
         *,
-        shape: ChunkCoords,
+        shape: ShapeLike,
         dtype: npt.DTypeLike = "float64",
         fill_value: Any | None = None,
         attributes: dict[str, JSON] | None = None,
@@ -380,7 +381,7 @@ async def create_array(
         codecs: Iterable[Codec | dict[str, JSON]] | None = None,
         dimension_names: Iterable[str] | None = None,
         # v2 only
-        chunks: ChunkCoords | None = None,
+        chunks: ShapeLike | None = None,
         dimension_separator: Literal[".", "/"] | None = None,
         order: Literal["C", "F"] | None = None,
         filters: list[dict[str, JSON]] | None = None,
@@ -890,7 +891,7 @@ def create_array(
         self,
         name: str,
         *,
-        shape: ChunkCoords,
+        shape: ShapeLike,
         dtype: npt.DTypeLike = "float64",
         fill_value: Any | None = None,
         attributes: dict[str, JSON] | None = None,
@@ -905,7 +906,7 @@ def create_array(
         codecs: Iterable[Codec | dict[str, JSON]] | None = None,
         dimension_names: Iterable[str] | None = None,
         # v2 only
-        chunks: ChunkCoords | None = None,
+        chunks: ShapeLike | None = None,
         dimension_separator: Literal[".", "/"] | None = None,
         order: Literal["C", "F"] | None = None,
         filters: list[dict[str, JSON]] | None = None,

tests/v3/package_with_entrypoint/__init__.py

@@ -2,7 +2,7 @@
 
 from numpy import ndarray
 
-from zarr.abc.codec import ArrayBytesCodec, CodecInput, CodecPipeline
+from zarr.abc.codec import ArrayBytesCodec, CodecInput, CodecOutput, CodecPipeline
 from zarr.codecs import BytesCodec
 from zarr.core.array_spec import ArraySpec
 from zarr.core.buffer import Buffer, NDBuffer
@@ -15,7 +15,7 @@ class TestEntrypointCodec(ArrayBytesCodec):
     async def encode(
         self,
         chunks_and_specs: Iterable[tuple[CodecInput | None, ArraySpec]],
-    ) -> BytesLike | None:
+    ) -> Iterable[CodecOutput | None]:
         pass
 
     async def decode(

tests/v3/test_api.py

@@ -1,12 +1,14 @@
+import pathlib
+
 import numpy as np
 import pytest
 from numpy.testing import assert_array_equal
-from pytest_asyncio import fixture
 
 import zarr
 from zarr import Array, Group
 from zarr.abc.store import Store
 from zarr.api.synchronous import create, load, open, open_group, save, save_array, save_group
+from zarr.store.memory import MemoryStore
 
 
 def test_create_array(memory_store: Store) -> None:
@@ -30,7 +32,7 @@ def test_create_array(memory_store: Store) -> None:
     assert z.chunks == (40,)
 
 
-async def test_open_array(memory_store: Store) -> None:
+async def test_open_array(memory_store: MemoryStore) -> None:
     store = memory_store
 
     # open array, create if doesn't exist
@@ -57,7 +59,7 @@ async def test_open_array(memory_store: Store) -> None:
         open(store="doesnotexist", mode="r")
 
 
-async def test_open_group(memory_store: Store) -> None:
+async def test_open_group(memory_store: MemoryStore) -> None:
     store = memory_store
 
     # open group, create if doesn't exist
@@ -85,59 +87,65 @@ def test_save_errors() -> None:
         save_group("data/group.zarr")
     with pytest.raises(TypeError):
         # no array provided
-        save_array("data/group.zarr")
+        save_array("data/group.zarr")  # type: ignore[call-arg]
     with pytest.raises(ValueError):
         # no arrays provided
         save("data/group.zarr")
 
 
-@fixture
-def tmppath(tmpdir):
-    return str(tmpdir / "example.zarr")
-
-
-def test_open_with_mode_r(tmppath) -> None:
+def test_open_with_mode_r(tmp_path: pathlib.Path) -> None:
     # 'r' means read only (must exist)
     with pytest.raises(FileNotFoundError):
-        zarr.open(store=tmppath, mode="r")
-    zarr.ones(store=tmppath, shape=(3, 3))
-    z2 = zarr.open(store=tmppath, mode="r")
+        zarr.open(store=tmp_path, mode="r")
+    zarr.ones(store=tmp_path, shape=(3, 3))
+    z2 = zarr.open(store=tmp_path, mode="r")
+    assert isinstance(z2, Array)
     assert (z2[:] == 1).all()
     with pytest.raises(ValueError):
         z2[:] = 3
 
 
-def test_open_with_mode_r_plus(tmppath) -> None:
+def test_open_with_mode_r_plus(tmp_path: pathlib.Path) -> None:
     # 'r+' means read/write (must exist)
     with pytest.raises(FileNotFoundError):
-        zarr.open(store=tmppath, mode="r+")
-    zarr.ones(store=tmppath, shape=(3, 3))
-    z2 = zarr.open(store=tmppath, mode="r+")
+        zarr.open(store=tmp_path, mode="r+")
+    zarr.ones(store=tmp_path, shape=(3, 3))
+    z2 = zarr.open(store=tmp_path, mode="r+")
+    assert isinstance(z2, Array)
     assert (z2[:] == 1).all()
     z2[:] = 3
 
 
-def test_open_with_mode_a(tmppath) -> None:
+def test_open_with_mode_a(tmp_path: pathlib.Path) -> None:
     # 'a' means read/write (create if doesn't exist)
-    zarr.open(store=tmppath, mode="a", shape=(3, 3))[...] = 1
-    z2 = zarr.open(store=tmppath, mode="a")
+    arr = zarr.open(store=tmp_path, mode="a", shape=(3, 3))
+    assert isinstance(arr, Array)
+    arr[...] = 1
+    z2 = zarr.open(store=tmp_path, mode="a")
+    assert isinstance(z2, Array)
     assert (z2[:] == 1).all()
     z2[:] = 3
 
 
-def test_open_with_mode_w(tmppath) -> None:
+def test_open_with_mode_w(tmp_path: pathlib.Path) -> None:
     # 'w' means create (overwrite if exists);
-    zarr.open(store=tmppath, mode="w", shape=(3, 3))[...] = 3
-    z2 = zarr.open(store=tmppath, mode="w", shape=(3, 3))
+    arr = zarr.open(store=tmp_path, mode="w", shape=(3, 3))
+    assert isinstance(arr, Array)
+
+    arr[...] = 3
+    z2 = zarr.open(store=tmp_path, mode="w", shape=(3, 3))
+    assert isinstance(z2, Array)
     assert not (z2[:] == 3).all()
     z2[:] = 3
 
 
-def test_open_with_mode_w_minus(tmppath) -> None:
+def test_open_with_mode_w_minus(tmp_path: pathlib.Path) -> None:
     # 'w-' means create  (fail if exists)
-    zarr.open(store=tmppath, mode="w-", shape=(3, 3))[...] = 1
+    arr = zarr.open(store=tmp_path, mode="w-", shape=(3, 3))
+    assert isinstance(arr, Array)
+    arr[...] = 1
     with pytest.raises(FileExistsError):
-        zarr.open(store=tmppath, mode="w-")
+        zarr.open(store=tmp_path, mode="w-")
 
 
 # def test_lazy_loader():

tests/v3/test_buffer.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import types
+
 import numpy as np
 import pytest
 
@@ -27,14 +29,14 @@
     cp = None
 
 
-def test_nd_array_like(xp):
+def test_nd_array_like(xp: types.ModuleType) -> None:
     ary = xp.arange(10)
     assert isinstance(ary, ArrayLike)
     assert isinstance(ary, NDArrayLike)
 
 
 @pytest.mark.asyncio
-async def test_async_array_prototype():
+async def test_async_array_prototype() -> None:
     """Test the use of a custom buffer prototype"""
 
     expect = np.zeros((9, 9), dtype="uint16", order="F")
@@ -55,13 +57,15 @@ async def test_async_array_prototype():
         prototype=my_prototype,
     )
     got = await a.getitem(selection=(slice(0, 9), slice(0, 9)), prototype=my_prototype)
-    assert isinstance(got, TestNDArrayLike)
-    assert np.array_equal(expect, got)
+    # ignoring a mypy error here that TestNDArrayLike doesn't meet the NDArrayLike protocol
+    # The test passes, so it clearly does.
+    assert isinstance(got, TestNDArrayLike)  # type: ignore[unreachable]
+    assert np.array_equal(expect, got)  # type: ignore[unreachable]
 
 
 @gpu_test
 @pytest.mark.asyncio
-async def test_async_array_gpu_prototype():
+async def test_async_array_gpu_prototype() -> None:
     """Test the use of the GPU buffer prototype"""
 
     expect = cp.zeros((9, 9), dtype="uint16", order="F")
@@ -85,7 +89,7 @@ async def test_async_array_gpu_prototype():
 
 
 @pytest.mark.asyncio
-async def test_codecs_use_of_prototype():
+async def test_codecs_use_of_prototype() -> None:
     expect = np.zeros((10, 10), dtype="uint16", order="F")
     a = await AsyncArray.create(
         StorePath(StoreExpectingTestBuffer(mode="w")) / "test_codecs_use_of_prototype",
@@ -112,13 +116,15 @@ async def test_codecs_use_of_prototype():
         prototype=my_prototype,
     )
     got = await a.getitem(selection=(slice(0, 10), slice(0, 10)), prototype=my_prototype)
-    assert isinstance(got, TestNDArrayLike)
-    assert np.array_equal(expect, got)
+    # ignoring a mypy error here that TestNDArrayLike doesn't meet the NDArrayLike protocol
+    # The test passes, so it clearly does.
+    assert isinstance(got, TestNDArrayLike)  # type: ignore[unreachable]
+    assert np.array_equal(expect, got)  # type: ignore[unreachable]
 
 
 @gpu_test
 @pytest.mark.asyncio
-async def test_codecs_use_of_gpu_prototype():
+async def test_codecs_use_of_gpu_prototype() -> None:
     expect = cp.zeros((10, 10), dtype="uint16", order="F")
     a = await AsyncArray.create(
         StorePath(MemoryStore(mode="w")) / "test_codecs_use_of_gpu_prototype",
@@ -147,7 +153,7 @@ async def test_codecs_use_of_gpu_prototype():
     assert cp.array_equal(expect, got)
 
 
-def test_numpy_buffer_prototype():
+def test_numpy_buffer_prototype() -> None:
     buffer = cpu.buffer_prototype.buffer.create_zero_length()
     ndbuffer = cpu.buffer_prototype.nd_buffer.create(shape=(1, 2), dtype=np.dtype("int64"))
     assert isinstance(buffer.as_array_like(), np.ndarray)