fix: some NDArrayLike fixes

Author: ilan-gold

src/zarr/_compat.py

@@ -102,7 +102,7 @@ def _reshape_view(arr: "NDArray[Any]", shape: tuple[int, ...]) -> "NDArray[Any]"
         If a view cannot be created (the array is not contiguous) on NumPy >= 2.1.
     """
     if Version(np.__version__) >= Version("2.1"):
-        return arr.reshape(shape, copy=False)  # type: ignore[call-overload, no-any-return]
+        return arr.reshape(shape, copy=False)
     else:
         arr.shape = shape
         return arr

src/zarr/core/buffer/core.py

@@ -71,7 +71,12 @@ def __setitem__(self, key: slice, value: Any) -> None: ...
     def __array__(self) -> npt.NDArray[Any]: ...
 
     def reshape(
-        self, shape: tuple[int, ...] | Literal[-1], *, order: Literal["A", "C", "F"] = ...
+        self,
+        shape: Sequence[SupportsIndex] | SupportsIndex,
+        /,
+        *,
+        order: Literal["A", "C", "F"] | None = ...,
+        copy: bool | None = ...,
     ) -> Self: ...
 
     def view(self, dtype: npt.DTypeLike) -> Self: ...
@@ -90,9 +95,16 @@ def copy(self) -> Self: ...
 
     def transpose(self, axes: SupportsIndex | Sequence[SupportsIndex] | None) -> Self: ...
 
-    def ravel(self, order: Literal["K", "A", "C", "F"] = ...) -> Self: ...
+    def ravel(self, order: Literal["K", "A", "C", "F"] | None = ...) -> Self: ...
 
-    def all(self) -> bool: ...
+    def all(
+        self,
+        axis: int | tuple[int, ...] | None = ...,
+        out: NDArrayLike | None = ...,
+        keepdims: Literal[False, 0] = ...,
+        *,
+        where: bool = ...,
+    ) -> bool: ...
 
     def __eq__(self, other: object) -> Self:  # type: ignore[override]
         """Element-wise equal

tests/test_api.py

@@ -2,7 +2,7 @@
 
 import inspect
 import re
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 import zarr.codecs
 import zarr.storage
@@ -558,7 +558,7 @@ def test_load_array(sync_store: Store) -> None:
 
     # can also load arrays directly into a numpy array
     for array_name in ["foo", "bar"]:
-        array = load(store, path=array_name)
+        array = cast(np.ndarray, load(store, path=array_name))
         assert isinstance(array, np.ndarray)
         if array_name == "foo":
             assert_array_equal(foo, array)
@@ -575,11 +575,11 @@ def test_load_zip(tmp_path: Path, path: str | None, load_read_only: bool | None)
     with ZipStore(file, mode="w", read_only=False) as zs:
         save(zs, data, path=path)
     with ZipStore(file, mode="r", read_only=load_read_only) as zs:
-        result = zarr.load(store=zs, path=path)
+        result = cast(np.ndarray, zarr.load(store=zs, path=path))
         assert isinstance(result, np.ndarray)
         assert np.array_equal(result, data)
     with ZipStore(file, read_only=load_read_only) as zs:
-        result = zarr.load(store=zs, path=path)
+        result = cast(np.ndarray, zarr.load(store=zs, path=path))
         assert isinstance(result, np.ndarray)
         assert np.array_equal(result, data)
 
@@ -593,7 +593,7 @@ def test_load_local(tmp_path: Path, path: str | None, load_read_only: bool) -> N
     with LocalStore(file, read_only=False) as zs:
         save(zs, data, path=path)
     with LocalStore(file, read_only=load_read_only) as zs:
-        result = zarr.load(store=zs, path=path)
+        result = cast(np.ndarray, zarr.load(store=zs, path=path))
         assert isinstance(result, np.ndarray)
         assert np.array_equal(result, data)
 

tests/test_codecs/test_codecs.py

@@ -2,7 +2,7 @@
 
 import json
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
 import pytest
@@ -145,7 +145,7 @@ def test_order_implicit(
 
     with config.set({"array.order": runtime_read_order}):
         a = Array.open(spath)
-    read_data = a[:, :]
+    read_data = cast(np.ndarray, a[:, :])
     assert np.array_equal(data, read_data)
 
     assert isinstance(read_data, np.ndarray)