Add numpy to mypy pre-commit check env

Author: dstansby

.pre-commit-config.yaml

@@ -26,8 +26,8 @@ repos:
     hooks:
       - id: mypy
         files: src
-        args: []
         additional_dependencies:
           - types-redis
           - types-setuptools
           - pytest
+          - numpy

src/zarr/array.py

@@ -417,7 +417,7 @@ async def setitem(
         # We accept any ndarray like object from the user and convert it
         # to a NDBuffer (or subclass). From this point onwards, we only pass
         # Buffer and NDBuffer between components.
-        value = factory(value)
+        value_buffer = factory(value)
 
         # merging with existing data and encoding chunks
         await self.metadata.codec_pipeline.write(
@@ -430,7 +430,7 @@ async def setitem(
                 )
                 for chunk_coords, chunk_selection, out_selection in indexer
             ],
-            value,
+            value_buffer,
         )
 
     async def resize(

src/zarr/buffer.py

@@ -11,6 +11,7 @@
 )
 
 import numpy as np
+import numpy.typing as npt
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -20,8 +21,8 @@
 
 # TODO: create a protocol for the attributes we need, for now we alias Numpy's ndarray
 #       both for the array-like and ndarray-like
-ArrayLike: TypeAlias = np.ndarray
-NDArrayLike: TypeAlias = np.ndarray
+ArrayLike: TypeAlias = npt.NDArray[Any]
+NDArrayLike: TypeAlias = npt.NDArray[Any]
 
 
 def check_item_key_is_1d_contiguous(key: Any) -> None:
@@ -40,7 +41,7 @@ def __call__(
             self,
             *,
             shape: Iterable[int],
-            dtype: np.DTypeLike,
+            dtype: npt.DTypeLike,
             order: Literal["C", "F"],
             fill_value: Any | None,
         ) -> NDBuffer:
@@ -163,7 +164,7 @@ def as_array_like(self) -> NDArrayLike:
         """
         return self._data
 
-    def as_nd_buffer(self, *, dtype: np.DTypeLike) -> NDBuffer:
+    def as_nd_buffer(self, *, dtype: npt.DTypeLike) -> NDBuffer:
         """Create a new NDBuffer from this one.
 
         This will never copy data.
@@ -179,7 +180,7 @@ def as_nd_buffer(self, *, dtype: np.DTypeLike) -> NDBuffer:
         """
         return NDBuffer.from_ndarray_like(self._data.view(dtype=dtype))
 
-    def as_numpy_array(self) -> np.ndarray:
+    def as_numpy_array(self) -> npt.NDArray[Any]:
         """Return the buffer as a NumPy array (host memory).
 
         Warning
@@ -271,7 +272,7 @@ def create(
         cls,
         *,
         shape: Iterable[int],
-        dtype: np.DTypeLike,
+        dtype: npt.DTypeLike,
         order: Literal["C", "F"] = "C",
         fill_value: Any | None = None,
     ) -> Self:
@@ -298,7 +299,7 @@ def create(
         A subclass can overwrite this method to create a ndarray-like object
         other then the default Numpy array.
         """
-        ret = cls(np.empty(shape=shape, dtype=dtype, order=order))
+        ret = cls(np.empty(shape=tuple(shape), dtype=dtype, order=order))
         if fill_value is not None:
             ret.fill(fill_value)
         return ret
@@ -319,7 +320,7 @@ def from_ndarray_like(cls, ndarray_like: NDArrayLike) -> Self:
         return cls(ndarray_like)
 
     @classmethod
-    def from_numpy_array(cls, array_like: np.ArrayLike) -> Self:
+    def from_numpy_array(cls, array_like: npt.ArrayLike) -> Self:
         """Create a new buffer of Numpy array-like object
 
         Parameters
@@ -360,7 +361,7 @@ def as_buffer(self) -> Buffer:
             data = np.ascontiguousarray(self._data)
         return Buffer(data.reshape(-1).view(dtype="b"))  # Flatten the array without copy
 
-    def as_numpy_array(self) -> np.ndarray:
+    def as_numpy_array(self) -> npt.NDArray[Any]:
         """Return the buffer as a NumPy array (host memory).
 
         Warning
@@ -393,9 +394,9 @@ def byteorder(self) -> Endian:
             return Endian(sys.byteorder)
 
     def reshape(self, newshape: Iterable[int]) -> Self:
-        return self.__class__(self._data.reshape(newshape))
+        return self.__class__(self._data.reshape(tuple(newshape)))
 
-    def astype(self, dtype: np.DTypeLike, order: Literal["K", "A", "C", "F"] = "K") -> Self:
+    def astype(self, dtype: npt.DTypeLike, order: Literal["K", "A", "C", "F"] = "K") -> Self:
         return self.__class__(self._data.astype(dtype=dtype, order=order))
 
     def __getitem__(self, key: Any) -> Self:
@@ -422,7 +423,7 @@ def transpose(self, *axes: np.SupportsIndex) -> Self:
         return self.__class__(self._data.transpose(*axes))
 
 
-def as_numpy_array_wrapper(func: Callable[[np.ndarray], bytes], buf: Buffer) -> Buffer:
+def as_numpy_array_wrapper(func: Callable[[npt.NDArray[Any]], bytes], buf: Buffer) -> Buffer:
     """Converts the input of `func` to a numpy array and the output back to `Buffer`.
 
     This function is useful when calling a `func` that only support host memory such

src/zarr/codecs/bytes.py

@@ -92,7 +92,9 @@ async def encode_single(
         assert isinstance(chunk_array, NDBuffer)
         if chunk_array.dtype.itemsize > 1:
             if self.endian is not None and self.endian != chunk_array.byteorder:
-                new_dtype = chunk_array.dtype.newbyteorder(self.endian.name)
+                # type-ignore is a numpy bug
+                # see https://github.com/numpy/numpy/issues/26473
+                new_dtype = chunk_array.dtype.newbyteorder(self.endian.name) # type: ignore[arg-type]
                 chunk_array = chunk_array.astype(new_dtype)
         return chunk_array.as_buffer()
 

src/zarr/codecs/sharding.py

@@ -8,6 +8,7 @@
 from typing import TYPE_CHECKING, NamedTuple
 
 import numpy as np
+import numpy.typing as npt
 
 from zarr.abc.codec import ByteGetter, ByteSetter, Codec, CodecPipeline
 from zarr.buffer import Buffer, NDBuffer
@@ -82,7 +83,7 @@ async def delete(self) -> None:
 
 class _ShardIndex(NamedTuple):
     # dtype uint64, shape (chunks_per_shard_0, chunks_per_shard_1, ..., 2)
-    offsets_and_lengths: np.ndarray
+    offsets_and_lengths: npt.NDArray[np.uint64]
 
     @property
     def chunks_per_shard(self) -> ChunkCoords:
@@ -97,7 +98,7 @@ def _localize_chunk(self, chunk_coords: ChunkCoords) -> ChunkCoords:
     def is_all_empty(self) -> bool:
         return bool(np.array_equiv(self.offsets_and_lengths, MAX_UINT_64))
 
-    def get_full_chunk_map(self) -> np.ndarray:
+    def get_full_chunk_map(self) -> npt.NDArray[np.bool_]:
         return self.offsets_and_lengths[..., 0] != MAX_UINT_64
 
     def get_chunk_slice(self, chunk_coords: ChunkCoords) -> tuple[int, int] | None:

src/zarr/metadata.py

@@ -399,7 +399,7 @@ def to_dict(self) -> JSON:
         zarray_dict["chunks"] = self.chunk_grid.chunk_shape
 
         _ = zarray_dict.pop("data_type")
-        zarray_dict["dtype"] = self.data_type
+        zarray_dict["dtype"] = self.data_type.str
 
         return zarray_dict
 

src/zarr/store/local.py

@@ -58,10 +58,10 @@ def _put(
     if start is not None:
         with path.open("r+b") as f:
             f.seek(start)
-            f.write(value.as_numpy_array())
+            f.write(value.as_numpy_array().tobytes())
         return None
     else:
-        return path.write_bytes(value.as_numpy_array())
+        return path.write_bytes(value.as_numpy_array().tobytes())
 
 
 class LocalStore(Store):

test.py

@@ -0,0 +1,3 @@
+import numpy as np
+a = np.array([1, 2, 3])
+a.dtype.newbyteorder("little")