diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3c0c4f9f92..ef5123b430 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index e7b5709b77..015a24c112 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -27,8 +27,10 @@ ZARRAY_JSON, ZATTRS_JSON, ChunkCoords, + ShapeLike, ZarrFormat, concurrent_map, + parse_shapelike, product, ) from zarr.core.config import config, parse_indexing_order @@ -113,7 +115,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, @@ -129,7 +131,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, @@ -140,9 +142,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.") @@ -214,7 +221,7 @@ async def _create_v3( cls, store_path: StorePath, *, - shape: ChunkCoords, + shape: ShapeLike, dtype: npt.DTypeLike, chunk_shape: ChunkCoords, fill_value: Any | None = None, @@ -232,6 +239,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: diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index f912872251..2a5479517e 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -16,6 +16,7 @@ JSON, ChunkCoords, ChunkCoordsLike, + ShapeLike, parse_named_configuration, parse_shapelike, ) @@ -26,7 +27,7 @@ def _guess_chunks( - shape: ChunkCoords, + shape: ShapeLike, typesize: int, *, increment_bytes: int = 256 * 1024, @@ -56,6 +57,8 @@ def _guess_chunks( ChunkCoords """ + if isinstance(shape, int): + shape = (shape,) ndims = len(shape) # require chunks to have non-zero length for all dimensions diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 8dbad2a496..01642a7511 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -25,6 +25,7 @@ ZATTRS_JSON, ZGROUP_JSON, ChunkCoords, + ShapeLike, ZarrFormat, parse_shapelike, ) @@ -366,7 +367,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, @@ -381,7 +382,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, @@ -891,7 +892,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, @@ -906,7 +907,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, diff --git a/tests/v3/package_with_entrypoint/__init__.py b/tests/v3/package_with_entrypoint/__init__.py index 0e84cde0d5..352b9d570e 100644 --- a/tests/v3/package_with_entrypoint/__init__.py +++ b/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( diff --git a/tests/v3/test_api.py b/tests/v3/test_api.py index 002678be88..239dd1c3e2 100644 --- a/tests/v3/test_api.py +++ b/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(): diff --git a/tests/v3/test_buffer.py b/tests/v3/test_buffer.py index 419f01c720..06c0714333 100644 --- a/tests/v3/test_buffer.py +++ b/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)