More typing fixes for tests (zarr-developers#2173)

Author: dstansby

pyproject.toml

@@ -241,7 +241,8 @@ ignore_errors = true
 module = [
     "tests.v2.*",
     "tests.v3.package_with_entrypoint.*",
-    "tests.v3.test_codecs.*",
+    "tests.v3.test_codecs.test_codecs",
+    "tests.v3.test_codecs.test_transpose",
     "tests.v3.test_metadata.*",
     "tests.v3.test_store.*",
     "tests.v3.test_config",

src/zarr/core/group.py

@@ -516,7 +516,7 @@ async def require_array(
         self,
         name: str,
         *,
-        shape: ChunkCoords,
+        shape: ShapeLike,
         dtype: npt.DTypeLike = None,
         exact: bool = False,
         **kwargs: Any,

tests/v3/package_with_entrypoint/__init__.py

@@ -2,6 +2,7 @@
 
 from numpy import ndarray
 
+import zarr.core.buffer
 from zarr.abc.codec import ArrayBytesCodec, CodecInput, CodecOutput, CodecPipeline
 from zarr.codecs import BytesCodec
 from zarr.core.array_spec import ArraySpec
@@ -29,7 +30,7 @@ def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) ->
 
 
 class TestEntrypointCodecPipeline(CodecPipeline):
-    def __init__(self, batch_size: int = 1):
+    def __init__(self, batch_size: int = 1) -> None:
         pass
 
     async def encode(
@@ -55,10 +56,10 @@ class TestEntrypointGroup:
     class Codec(BytesCodec):
         pass
 
-    class Buffer(Buffer):
+    class Buffer(zarr.core.buffer.Buffer):
         pass
 
-    class NDBuffer(NDBuffer):
+    class NDBuffer(zarr.core.buffer.NDBuffer):
         pass
 
     class Pipeline(CodecPipeline):

tests/v3/test_codecs/test_blosc.py

@@ -24,10 +24,9 @@ async def test_blosc_evolve(store: Store, dtype: str) -> None:
         fill_value=0,
         codecs=[BytesCodec(), BloscCodec()],
     )
-
-    zarr_json = json.loads(
-        (await store.get(f"{path}/zarr.json", prototype=default_buffer_prototype())).to_bytes()
-    )
+    buf = await store.get(f"{path}/zarr.json", prototype=default_buffer_prototype())
+    assert buf is not None
+    zarr_json = json.loads(buf.to_bytes())
     blosc_configuration_json = zarr_json["codecs"][1]["configuration"]
     assert blosc_configuration_json["typesize"] == typesize
     if typesize == 1:
@@ -45,10 +44,9 @@ async def test_blosc_evolve(store: Store, dtype: str) -> None:
         fill_value=0,
         codecs=[ShardingCodec(chunk_shape=(16, 16), codecs=[BytesCodec(), BloscCodec()])],
     )
-
-    zarr_json = json.loads(
-        (await store.get(f"{path2}/zarr.json", prototype=default_buffer_prototype())).to_bytes()
-    )
+    buf = await store.get(f"{path2}/zarr.json", prototype=default_buffer_prototype())
+    assert buf is not None
+    zarr_json = json.loads(buf.to_bytes())
     blosc_configuration_json = zarr_json["codecs"][0]["configuration"]["codecs"][1]["configuration"]
     assert blosc_configuration_json["typesize"] == typesize
     if typesize == 1:

tests/v3/test_codecs/test_codecs.py

@@ -7,7 +7,7 @@
 import numpy as np
 import pytest
 
-import zarr.v2
+import zarr.v2.creation
 from zarr import Array, AsyncArray, config
 from zarr.codecs import (
     BytesCodec,
@@ -23,6 +23,7 @@
 if TYPE_CHECKING:
     from zarr.abc.codec import Codec
     from zarr.abc.store import Store
+    from zarr.core.buffer.core import NDArrayLike
     from zarr.core.common import MemoryOrder
 
 
@@ -39,7 +40,7 @@ class _AsyncArraySelectionProxy:
     array: AsyncArray
     selection: Selection
 
-    async def get(self) -> np.ndarray:
+    async def get(self) -> NDArrayLike:
         return await self.array.getitem(self.selection)
 
     async def set(self, value: np.ndarray) -> None:
@@ -119,7 +120,7 @@ async def test_order(
 
     if not with_sharding:
         # Compare with zarr-python
-        z = zarr.v2.create(
+        z = zarr.v2.creation.create(
             shape=data.shape,
             chunks=(32, 8),
             dtype="<u2",
@@ -268,7 +269,7 @@ async def test_zarr_compat(store: Store) -> None:
         fill_value=1,
     )
 
-    z2 = zarr.v2.create(
+    z2 = zarr.v2.creation.create(
         shape=data.shape,
         chunks=(10, 10),
         dtype=data.dtype,
@@ -310,7 +311,7 @@ async def test_zarr_compat_F(store: Store) -> None:
         codecs=[TransposeCodec(order=order_from_dim("F", data.ndim)), BytesCodec()],
     )
 
-    z2 = zarr.v2.create(
+    z2 = zarr.v2.creation.create(
         shape=data.shape,
         chunks=(10, 10),
         dtype=data.dtype,
@@ -406,7 +407,7 @@ def test_invalid_metadata(store: Store) -> None:
             fill_value=0,
             codecs=[
                 BytesCodec(),
-                TransposeCodec(order="F"),
+                TransposeCodec(order="F"),  # type: ignore[arg-type]
             ],
         )
     spath4 = StorePath(store, "invalid_missing_bytes_codec")

tests/v3/test_codecs/test_endian.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-import zarr.v2
+import zarr.v2.creation
 from zarr import AsyncArray
 from zarr.abc.store import Store
 from zarr.codecs import BytesCodec
@@ -35,7 +35,7 @@ async def test_endian(store: Store, endian: Literal["big", "little"]) -> None:
     assert np.array_equal(data, readback_data)
 
     # Compare with v2
-    z = zarr.v2.create(
+    z = zarr.v2.creation.create(
         shape=data.shape,
         chunks=(16, 16),
         dtype=">u2" if endian == "big" else "<u2",
@@ -74,7 +74,7 @@ async def test_endian_write(
     assert np.array_equal(data, readback_data)
 
     # Compare with zarr-python
-    z = zarr.v2.create(
+    z = zarr.v2.creation.create(
         shape=data.shape,
         chunks=(16, 16),
         dtype=">u2" if dtype_store_endian == "big" else "<u2",

tests/v3/test_codecs/test_sharding.py

@@ -1,6 +1,8 @@
 import pickle
+from typing import Any
 
 import numpy as np
+import numpy.typing as npt
 import pytest
 
 from zarr import Array, AsyncArray
@@ -32,7 +34,10 @@
 )
 @pytest.mark.parametrize("offset", [0, 10])
 def test_sharding(
-    store: Store, array_fixture: np.ndarray, index_location: ShardingCodecIndexLocation, offset: int
+    store: Store,
+    array_fixture: npt.NDArray[Any],
+    index_location: ShardingCodecIndexLocation,
+    offset: int,
 ) -> None:
     """
     Test that we can create an array with a sharding codec, write data to that array, and get
@@ -80,7 +85,7 @@ def test_sharding(
     indirect=["array_fixture"],
 )
 def test_sharding_partial(
-    store: Store, array_fixture: np.ndarray, index_location: ShardingCodecIndexLocation
+    store: Store, array_fixture: npt.NDArray[Any], index_location: ShardingCodecIndexLocation
 ) -> None:
     data = array_fixture
     spath = StorePath(store)
@@ -123,7 +128,7 @@ def test_sharding_partial(
 @pytest.mark.parametrize("index_location", ["start", "end"])
 @pytest.mark.parametrize("store", ("local", "memory"), indirect=["store"])
 def test_sharding_partial_read(
-    store: Store, array_fixture: np.ndarray, index_location: ShardingCodecIndexLocation
+    store: Store, array_fixture: npt.NDArray[Any], index_location: ShardingCodecIndexLocation
 ) -> None:
     data = array_fixture
     spath = StorePath(store)
@@ -160,7 +165,7 @@ def test_sharding_partial_read(
 @pytest.mark.parametrize("index_location", ["start", "end"])
 @pytest.mark.parametrize("store", ("local", "memory"), indirect=["store"])
 def test_sharding_partial_overwrite(
-    store: Store, array_fixture: np.ndarray, index_location: ShardingCodecIndexLocation
+    store: Store, array_fixture: npt.NDArray[Any], index_location: ShardingCodecIndexLocation
 ) -> None:
     data = array_fixture[:10, :10, :10]
     spath = StorePath(store)
@@ -212,7 +217,7 @@ def test_sharding_partial_overwrite(
 @pytest.mark.parametrize("store", ("local", "memory"), indirect=["store"])
 def test_nested_sharding(
     store: Store,
-    array_fixture: np.ndarray,
+    array_fixture: npt.NDArray[Any],
     outer_index_location: ShardingCodecIndexLocation,
     inner_index_location: ShardingCodecIndexLocation,
 ) -> None:

tests/v3/test_codecs/test_transpose.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-import zarr.v2
+import zarr.v2.creation
 from zarr import Array, AsyncArray, config
 from zarr.abc.store import Store
 from zarr.codecs import BytesCodec, ShardingCodec, TransposeCodec
@@ -72,7 +72,7 @@ async def test_transpose(
 
     if not with_sharding:
         # Compare with zarr-python
-        z = zarr.v2.create(
+        z = zarr.v2.creation.create(
             shape=data.shape,
             chunks=(1, 32, 8),
             dtype="<u2",

tests/v3/test_group.py

@@ -9,7 +9,7 @@
 from zarr import Array, AsyncArray, AsyncGroup, Group
 from zarr.api.synchronous import open_group
 from zarr.core.buffer import default_buffer_prototype
-from zarr.core.common import ZarrFormat
+from zarr.core.common import JSON, ZarrFormat
 from zarr.core.group import GroupMetadata
 from zarr.core.sync import sync
 from zarr.errors import ContainsArrayError, ContainsGroupError
@@ -409,7 +409,7 @@ def test_group_creation_existing_node(
     spath = StorePath(store)
     group = Group.create(spath, zarr_format=zarr_format)
     expected_exception: type[ContainsArrayError] | type[ContainsGroupError]
-    attributes = {"old": True}
+    attributes: dict[str, JSON] = {"old": True}
 
     if extant_node == "array":
         expected_exception = ContainsArrayError
@@ -645,7 +645,7 @@ async def test_asyncgroup_create_array(
     shape = (10,)
     dtype = "uint8"
     chunk_shape = (4,)
-    attributes = {"foo": 100}
+    attributes: dict[str, JSON] = {"foo": 100}
 
     sub_node_path = "sub_array"
     subnode = await agroup.create_array(