Merge branch 'v3' into upath

Author: d-v-b

.github/workflows/test.yml

@@ -21,8 +21,8 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.10', '3.11']
-        numpy-version: ['1.24', '1.26', '2.0.0rc1']
+        python-version: ['3.10', '3.11', '3.12']
+        numpy-version: ['1.24', '1.26', '2.0']
         dependency-set: ["minimal", "optional"]
 
     steps:

.pre-commit-config.yaml

@@ -7,7 +7,7 @@ default_language_version:
   python: python3
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.4.9'
+    rev: 'v0.5.0'
     hooks:
     - id: ruff
       args: ["--fix", "--show-fixes"]
@@ -22,7 +22,7 @@ repos:
     hooks:
     - id: check-yaml
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.10.0
+    rev: v1.10.1
     hooks:
       - id: mypy
         files: src

pyproject.toml

@@ -126,13 +126,13 @@ extra-dependencies = [
 features = ["extra"]
 
 [[tool.hatch.envs.test.matrix]]
-python = ["3.10", "3.11"]
-numpy = ["1.24", "1.26", "2.0.0rc1"]
+python = ["3.10", "3.11", "3.12"]
+numpy = ["1.24", "1.26", "2.0"]
 version = ["minimal"]
 
 [[tool.hatch.envs.test.matrix]]
-python = ["3.10", "3.11"]
-numpy = ["1.24", "1.26", "2.0.0rc1"]
+python = ["3.10", "3.11", "3.12"]
+numpy = ["1.24", "1.26", "2.0"]
 features = ["optional"]
 
 [tool.hatch.envs.test.scripts]
@@ -184,10 +184,7 @@ extend-select = [
     "RUF",
 ]
 ignore = [
-    "RUF003",
     "RUF005",
-    "RUF012",
-    "RUF015",
 ]
 
 [tool.mypy]

src/zarr/abc/codec.py

@@ -187,10 +187,7 @@ async def decode_partial(
         Iterable[NDBuffer | None]
         """
         return await concurrent_map(
-            [
-                (byte_getter, selection, chunk_spec)
-                for byte_getter, selection, chunk_spec in batch_info
-            ],
+            list(batch_info),
             self._decode_partial_single,
             config.get("async.concurrency"),
         )
@@ -227,10 +224,7 @@ async def encode_partial(
             The chunk spec contains information about the chunk.
         """
         await concurrent_map(
-            [
-                (byte_setter, chunk_array, selection, chunk_spec)
-                for byte_setter, chunk_array, selection, chunk_spec in batch_info
-            ],
+            list(batch_info),
             self._encode_partial_single,
             config.get("async.concurrency"),
         )
@@ -402,7 +396,7 @@ async def batching_helper(
     batch_info: Iterable[tuple[CodecInput | None, ArraySpec]],
 ) -> list[CodecOutput | None]:
     return await concurrent_map(
-        [(chunk_array, chunk_spec) for chunk_array, chunk_spec in batch_info],
+        list(batch_info),
         noop_for_none(func),
         config.get("async.concurrency"),
     )

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:

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

src/zarr/store/remote.py

@@ -69,7 +69,7 @@ def __init__(
             self.path = url.path.rstrip("/")
             self._fs = url.fs
         else:
-            raise ValueError("URL not understood, %s", url)
+            raise ValueError(f"URL not understood, {url}")
         self.allowed_exceptions = allowed_exceptions
         # test instantiate file system
         if not self._fs.async_impl:

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=}")