Fix indexing with bools

src/zarr/array.py

@@ -389,7 +389,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:

src/zarr/indexing.py

@@ -101,11 +101,13 @@ 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[int]:
+    """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]]:
@@ -118,6 +120,16 @@ def is_integer_list(x: Any) -> TypeGuard[list[int]]:
     return isinstance(x, list) and len(x) > 0 and is_integer(x[0])
 
 
+def is_bool_list(x: Any) -> TypeGuard[list[int]]:
+    """True if x is a list of boolean.
+
+    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_bool(x[0])
+
+
 def is_integer_array(x: Any, ndim: int | None = None) -> TypeGuard[npt.NDArray[np.intp]]:
     t = not np.isscalar(x) and hasattr(x, "shape") and hasattr(x, "dtype") and x.dtype.kind in "ui"
     if ndim is not None:
@@ -132,6 +144,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
@@ -143,7 +159,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
     ----------
@@ -156,9 +172,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 = (
@@ -180,19 +201,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
         )
     )
 

tests/v3/test_indexing.py

@@ -1264,6 +1264,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 = [
@@ -1336,6 +1338,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):
@@ -1718,3 +1723,49 @@ 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))
+    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)]
+    actual = z[*selection]
+    assert_array_equal(expected, actual, err_msg=f"{selection=}")