diff --git a/src/zarr/array.py b/src/zarr/array.py index 4318050dd5..26a19e64ab 100644 --- a/src/zarr/array.py +++ b/src/zarr/array.py @@ -403,7 +403,7 @@ def attrs(self) -> dict[str, JSON]: @property def read_only(self) -> bool: - return bool(~self.store_path.store.writeable) + return bool(not self.store_path.store.writeable) @property def path(self) -> str: diff --git a/src/zarr/indexing.py b/src/zarr/indexing.py index 29b6dd790a..ae4aa0681b 100644 --- a/src/zarr/indexing.py +++ b/src/zarr/indexing.py @@ -87,21 +87,23 @@ def ceildiv(a: float, b: float) -> int: def is_integer(x: Any) -> TypeGuard[int]: - """True if x is an integer (both pure Python or NumPy). + """True if x is an integer (both pure Python or NumPy).""" + return isinstance(x, numbers.Integral) and not is_bool(x) - Note that Python's bool is considered an integer too. - """ - return isinstance(x, numbers.Integral) + +def is_bool(x: Any) -> TypeGuard[bool | np.bool_]: + """True if x is a boolean (both pure Python or NumPy).""" + return type(x) in [bool, np.bool_] def is_integer_list(x: Any) -> TypeGuard[list[int]]: - """True if x is a list of integers. + """True if x is a list of integers.""" + return isinstance(x, list) and len(x) > 0 and all(is_integer(i) for i in x) - This function assumes ie *does not check* that all elements of the list - have the same type. Mixed type lists will result in other errors that will - bubble up anyway. - """ - return isinstance(x, list) and len(x) > 0 and is_integer(x[0]) + +def is_bool_list(x: Any) -> TypeGuard[list[bool | np.bool_]]: + """True if x is a list of boolean.""" + return isinstance(x, list) and len(x) > 0 and all(is_bool(i) for i in x) def is_integer_array(x: Any, ndim: int | None = None) -> TypeGuard[npt.NDArray[np.intp]]: @@ -118,6 +120,10 @@ def is_bool_array(x: Any, ndim: int | None = None) -> TypeGuard[npt.NDArray[np.b return t +def is_int_or_bool_iterable(x: Any) -> bool: + return is_integer_list(x) or is_integer_array(x) or is_bool_array(x) or is_bool_list(x) + + def is_scalar(value: Any, dtype: np.dtype[Any]) -> bool: if np.isscalar(value): return True @@ -129,7 +135,7 @@ def is_scalar(value: Any, dtype: np.dtype[Any]) -> bool: def is_pure_fancy_indexing(selection: Any, ndim: int) -> bool: - """Check whether a selection contains only scalars or integer array-likes. + """Check whether a selection contains only scalars or integer/bool array-likes. Parameters ---------- @@ -142,9 +148,14 @@ def is_pure_fancy_indexing(selection: Any, ndim: int) -> bool: True if the selection is a pure fancy indexing expression (ie not mixed with boolean or slices). """ + if is_bool_array(selection): + # is mask selection + return True + if ndim == 1: - if is_integer_list(selection) or is_integer_array(selection): + if is_integer_list(selection) or is_integer_array(selection) or is_bool_list(selection): return True + # if not, we go through the normal path below, because a 1-tuple # of integers is also allowed. no_slicing = ( @@ -166,19 +177,21 @@ def is_pure_orthogonal_indexing(selection: Selection, ndim: int) -> TypeGuard[Or if not ndim: return False - # Case 1: Selection is a single iterable of integers - if is_integer_list(selection) or is_integer_array(selection, ndim=1): + selection_normalized = (selection,) if not isinstance(selection, tuple) else selection + + # Case 1: Selection contains of iterable of integers or boolean + if len(selection_normalized) == ndim and all( + is_int_or_bool_iterable(s) for s in selection_normalized + ): return True - # Case two: selection contains either zero or one integer iterables. + # Case 2: selection contains either zero or one integer iterables. # All other selection elements are slices or integers return ( - isinstance(selection, tuple) - and len(selection) == ndim - and sum(is_integer_list(elem) or is_integer_array(elem) for elem in selection) <= 1 + len(selection_normalized) <= ndim + and sum(is_int_or_bool_iterable(s) for s in selection_normalized) <= 1 and all( - is_integer_list(elem) or is_integer_array(elem) or isinstance(elem, int | slice) - for elem in selection + is_int_or_bool_iterable(s) or isinstance(s, int | slice) for s in selection_normalized ) ) @@ -1023,7 +1036,7 @@ def __init__(self, selection: CoordinateSelection, shape: ChunkCoords, chunk_gri # flatten selection selection_broadcast = tuple(dim_sel.reshape(-1) for dim_sel in selection_broadcast) chunks_multi_index_broadcast = tuple( - dim_chunks.reshape(-1) for dim_chunks in chunks_multi_index_broadcast + [dim_chunks.reshape(-1) for dim_chunks in chunks_multi_index_broadcast] ) # ravel chunk indices diff --git a/tests/v3/test_indexing.py b/tests/v3/test_indexing.py index 13a7d953e1..c84c091089 100644 --- a/tests/v3/test_indexing.py +++ b/tests/v3/test_indexing.py @@ -204,7 +204,6 @@ def test_get_basic_selection_0d(store: StorePath, use_out: bool, value: Any, dty slice(50, 150, 10), ] - basic_selections_1d_bad = [ # only positive step supported slice(None, None, -1), @@ -305,7 +304,6 @@ def test_get_basic_selection_1d(store: StorePath): (Ellipsis, slice(None), slice(None)), ] - basic_selections_2d_bad = [ # bad stuff 2.3, @@ -1272,6 +1270,8 @@ def _test_get_mask_selection(a, z, selection): assert_array_equal(expect, actual) actual = z.vindex[selection] assert_array_equal(expect, actual) + actual = z[selection] + assert_array_equal(expect, actual) mask_selections_1d_bad = [ @@ -1344,6 +1344,9 @@ def _test_set_mask_selection(v, a, z, selection): z[:] = 0 z.vindex[selection] = v[selection] assert_array_equal(a, z[:]) + z[:] = 0 + z[selection] = v[selection] + assert_array_equal(a, z[:]) def test_set_mask_selection_1d(store: StorePath): @@ -1726,3 +1729,51 @@ def test_accessed_chunks(shape, chunks, ops): ) == 1 # Check that no other chunks were accessed assert len(delta_counts) == 0 + + +@pytest.mark.parametrize( + "selection", + [ + # basic selection + [...], + [1, ...], + [slice(None)], + [1, 3], + [[1, 2, 3], 9], + [np.arange(1000)], + [slice(5, 15)], + [slice(2, 4), 4], + [[1, 3]], + # mask selection + [np.tile([True, False], (1000, 5))], + [np.full((1000, 10), False)], + # coordinate selection + [[1, 2, 3, 4], [5, 6, 7, 8]], + [[100, 200, 300], [4, 5, 6]], + ], +) +def test_indexing_equals_numpy(store, selection): + a = np.arange(10000, dtype=int).reshape(1000, 10) + z = zarr_array_from_numpy_array(store, a, chunk_shape=(300, 3)) + # note: in python 3.10 a[*selection] is not valid unpacking syntax + expected = a[(*selection,)] + actual = z[(*selection,)] + assert_array_equal(expected, actual, err_msg=f"selection: {selection}") + + +@pytest.mark.parametrize( + "selection", + [ + [np.tile([True, False], 500), np.tile([True, False], 5)], + [np.full(1000, False), np.tile([True, False], 5)], + [np.full(1000, True), np.full(10, True)], + [np.full(1000, True), [True, False] * 5], + ], +) +def test_orthogonal_bool_indexing_like_numpy_ix(store, selection): + a = np.arange(10000, dtype=int).reshape(1000, 10) + z = zarr_array_from_numpy_array(store, a, chunk_shape=(300, 3)) + expected = a[np.ix_(*selection)] + # note: in python 3.10 z[*selection] is not valid unpacking syntax + actual = z[(*selection,)] + assert_array_equal(expected, actual, err_msg=f"{selection=}")