mypy fixes for NumPy 2.1.0 (#2139)

* mypy fixes for NumPy 2.1.0

This includes some fixes to get mypy type checks passing.

* fixup

* Update src/zarr/codecs/crc32c_.py

Co-authored-by: David Stansby <[email protected]>

---------

Co-authored-by: David Stansby <[email protected]>

Author: TomAugspurger

.pre-commit-config.yaml

@@ -22,7 +22,7 @@ repos:
     hooks:
     - id: check-yaml
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.11.1
+    rev: v1.11.2
     hooks:
       - id: mypy
         files: src

src/zarr/codecs/crc32c_.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
+import typing_extensions
 from crc32c import crc32c
 
 from zarr.abc.codec import BytesBytesCodec
@@ -37,7 +38,8 @@ async def _decode_single(
         crc32_bytes = data[-4:]
         inner_bytes = data[:-4]
 
-        computed_checksum = np.uint32(crc32c(inner_bytes)).tobytes()
+        # Need to do a manual cast until https://github.com/numpy/numpy/issues/26783 is resolved
+        computed_checksum = np.uint32(crc32c(cast(typing_extensions.Buffer, inner_bytes))).tobytes()
         stored_checksum = bytes(crc32_bytes)
         if computed_checksum != stored_checksum:
             raise ValueError(
@@ -52,7 +54,7 @@ async def _encode_single(
     ) -> Buffer | None:
         data = chunk_bytes.as_numpy_array()
         # Calculate the checksum and "cast" it to a numpy array
-        checksum = np.array([crc32c(data)], dtype=np.uint32)
+        checksum = np.array([crc32c(cast(typing_extensions.Buffer, data))], dtype=np.uint32)
         # Append the checksum (as bytes) to the data
         return chunk_spec.prototype.buffer.from_array_like(np.append(data, checksum.view("b")))
 

src/zarr/codecs/sharding.py

@@ -5,7 +5,7 @@
 from enum import Enum
 from functools import lru_cache
 from operator import itemgetter
-from typing import TYPE_CHECKING, Any, NamedTuple
+from typing import TYPE_CHECKING, Any, NamedTuple, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -101,7 +101,9 @@ class _ShardIndex(NamedTuple):
 
     @property
     def chunks_per_shard(self) -> ChunkCoords:
-        return self.offsets_and_lengths.shape[0:-1]
+        result = tuple(self.offsets_and_lengths.shape[0:-1])
+        # The cast is required until https://github.com/numpy/numpy/pull/27211 is merged
+        return cast(ChunkCoords, result)
 
     def _localize_chunk(self, chunk_coords: ChunkCoords) -> ChunkCoords:
         return tuple(

src/zarr/core/metadata.py

@@ -602,9 +602,21 @@ def parse_fill_value_v3(fill_value: Any, dtype: FLOAT_DTYPE) -> FLOAT: ...
 def parse_fill_value_v3(fill_value: Any, dtype: COMPLEX_DTYPE) -> COMPLEX: ...
 
 
+@overload
+def parse_fill_value_v3(fill_value: Any, dtype: np.dtype[Any]) -> Any:
+    # This dtype[Any] is unfortunately necessary right now.
+    # See https://github.com/zarr-developers/zarr-python/issues/2131#issuecomment-2318010899
+    # for more details, but `dtype` here (which comes from `parse_dtype`)
+    # is np.dtype[Any].
+    #
+    # If you want the specialized types rather than Any, you need to use `np.dtypes.<dtype>`
+    # rather than np.dtypes(<type>)
+    ...
+
+
 def parse_fill_value_v3(
-    fill_value: Any, dtype: BOOL_DTYPE | INTEGER_DTYPE | FLOAT_DTYPE | COMPLEX_DTYPE
-) -> BOOL | INTEGER | FLOAT | COMPLEX:
+    fill_value: Any, dtype: BOOL_DTYPE | INTEGER_DTYPE | FLOAT_DTYPE | COMPLEX_DTYPE | np.dtype[Any]
+) -> BOOL | INTEGER | FLOAT | COMPLEX | Any:
     """
     Parse `fill_value`, a potential fill value, into an instance of `dtype`, a data type.
     If `fill_value` is `None`, then this function will return the result of casting the value 0