factor array element iteration routines into stand-alone functions, and add a failing test

Author: d-v-b

src/zarr/core/array.py

@@ -4928,3 +4928,96 @@ def _parse_data_params(
             raise ValueError(msg)
         dtype_out = data.dtype
     return data, shape_out, dtype_out
+
+
+def iter_chunk_coords(
+    array: Array | AsyncArray[Any],
+    *,
+    origin: Sequence[int] | None = None,
+    selection_shape: Sequence[int] | None = None,
+) -> Iterator[ChunkCoords]:
+    """
+    Create an iterator over the coordinates of chunks in chunk grid space. If the `origin`
+    keyword is used, iteration will start at the chunk index specified by `origin`.
+    The default behavior is to start at the origin of the grid coordinate space.
+    If the `selection_shape` keyword is used, iteration will be bounded over a contiguous region
+    ranging from `[origin, origin selection_shape]`, where the upper bound is exclusive as
+    per python indexing conventions.
+
+    Parameters
+    ----------
+    array : Array | AsyncArray
+        The array to iterate over.
+    origin : Sequence[int] | None, default=None
+        The origin of the selection relative to the array's chunk grid.
+    selection_shape : Sequence[int] | None, default=None
+        The shape of the selection in chunk grid coordinates.
+
+    Yields
+    ------
+    chunk_coords: ChunkCoords
+        The coordinates of each chunk in the selection.
+    """
+    return _iter_grid(array.cdata_shape, origin=origin, selection_shape=selection_shape)
+
+
+def iter_chunk_keys(
+    array: Array | AsyncArray[Any],
+    *,
+    origin: Sequence[int] | None = None,
+    selection_shape: Sequence[int] | None = None,
+) -> Iterator[str]:
+    """
+    Iterate over the storage keys of each chunk, relative to an optional origin, and optionally
+    limited to a contiguous region in chunk grid coordinates.
+
+    Parameters
+    ----------
+    array : Array | AsyncArray
+        The array to iterate over.
+    origin : Sequence[int] | None, default=None
+        The origin of the selection relative to the array's chunk grid.
+    selection_shape : Sequence[int] | None, default=None
+        The shape of the selection in chunk grid coordinates.
+
+    Yields
+    ------
+    key: str
+        The storage key of each chunk in the selection.
+    """
+    # Iterate over the coordinates of chunks in chunk grid space.
+    for k in iter_chunk_coords(array, origin=origin, selection_shape=selection_shape):
+        # Encode the chunk key from the chunk coordinates.
+        yield array.metadata.encode_chunk_key(k)
+
+
+def iter_chunk_regions(
+    array: Array | AsyncArray[Any],
+    *,
+    origin: Sequence[int] | None = None,
+    selection_shape: Sequence[int] | None = None,
+) -> Iterator[tuple[slice, ...]]:
+    """
+    Iterate over the regions spanned by each chunk.
+
+    Parameters
+    ----------
+    array : Array | AsyncArray
+        The array to iterate over.
+    origin : Sequence[int] | None, default=None
+        The origin of the selection relative to the array's chunk grid.
+    selection_shape : Sequence[int] | None, default=None
+        The shape of the selection in chunk grid coordinates.
+
+    Yields
+    ------
+    region: tuple[slice, ...]
+        A tuple of slice objects representing the region spanned by each chunk in the selection.
+    """
+    for cgrid_position in iter_chunk_coords(array, origin=origin, selection_shape=selection_shape):
+        out: tuple[slice, ...] = ()
+        for c_pos, c_shape in zip(cgrid_position, array.chunks, strict=False):
+            start = c_pos * c_shape
+            stop = start + c_shape
+            out += (slice(start, stop, 1),)
+        yield out

tests/test_array.py

@@ -6,7 +6,7 @@
 import pickle
 import re
 import sys
-from itertools import accumulate
+from itertools import accumulate, starmap
 from typing import TYPE_CHECKING, Any, Literal
 from unittest import mock
 
@@ -37,6 +37,7 @@
     create_array,
     default_filters_v2,
     default_serializer_v3,
+    iter_chunk_keys,
 )
 from zarr.core.buffer import NDArrayLike, NDArrayLikeOrScalar, default_buffer_prototype
 from zarr.core.chunk_grids import _auto_partition
@@ -59,7 +60,7 @@
 from zarr.core.dtype.npy.common import NUMPY_ENDIANNESS_STR, endianness_from_numpy_str
 from zarr.core.dtype.npy.string import UTF8Base
 from zarr.core.group import AsyncGroup
-from zarr.core.indexing import BasicIndexer, ceildiv
+from zarr.core.indexing import BasicIndexer, _iter_grid, ceildiv
 from zarr.core.metadata.v2 import ArrayV2Metadata
 from zarr.core.metadata.v3 import ArrayV3Metadata
 from zarr.core.sync import sync
@@ -1835,3 +1836,35 @@ def test_unknown_object_codec_default_filters_v2() -> None:
     msg = f"Data type {dtype} requires an unknown object codec: {dtype.object_codec_id!r}."
     with pytest.raises(ValueError, match=re.escape(msg)):
         default_filters_v2(dtype)
+
+
+@pytest.mark.parametrize(
+    ("shard_size", "chunk_size"),
+    [
+        ((8,), (8,)),
+        ((8,), (2,)),
+        (
+            (
+                8,
+                10,
+            ),
+            (2, 2),
+        ),
+    ],
+)
+def test_iter_chunk_keys(shard_size: tuple[int, ...], chunk_size: tuple[int, ...]) -> None:
+    store = {}
+    arr = zarr.create_array(
+        store,
+        dtype="uint8",
+        shape=tuple(2 * x for x in shard_size),
+        chunks=chunk_size,
+        shards=shard_size,
+        zarr_format=3,
+    )
+    shard_grid_shape = tuple(starmap(ceildiv, zip(arr.shape, arr.shards, strict=True)))
+    expected_keys = tuple(
+        arr.metadata.chunk_key_encoding.encode_chunk_key(region)
+        for region in _iter_grid(shard_grid_shape)
+    )
+    assert tuple(iter_chunk_keys(arr)) == expected_keys