more progress on typing; still not passing mypy

Author: rabernat

src/zarr/core/metadata/v3.py

@@ -335,30 +335,16 @@ def update_attributes(self, attributes: dict[str, JSON]) -> Self:
 # enum Literals can't be used in typing, so we have to restate all of the V3 dtypes as types
 # https://github.com/python/typing/issues/781
 
-BOOL = np.bool_
-# BOOL_DTYPE = np.dtypes.BoolDType
 BOOL_DTYPE = Literal["bool"]
+BOOL = np.bool_
 INTEGER_DTYPE = Literal["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"]
-# INTEGER_DTYPE = (
-#     np.dtypes.Int8DType
-#     | np.dtypes.Int16DType
-#     | np.dtypes.Int32DType
-#     | np.dtypes.Int64DType
-#     | np.dtypes.UInt8DType
-#     | np.dtypes.UInt16DType
-#     | np.dtypes.UInt32DType
-#     | np.dtypes.UInt64DType
-# )
 INTEGER = np.int8 | np.int16 | np.int32 | np.int64 | np.uint8 | np.uint16 | np.uint32 | np.uint64
-# FLOAT_DTYPE = np.dtypes.Float16DType | np.dtypes.Float32DType | np.dtypes.Float64DType
 FLOAT_DTYPE = Literal["float16", "float32", "float64"]
 FLOAT = np.float16 | np.float32 | np.float64
-# COMPLEX_DTYPE = np.dtypes.Complex64DType | np.dtypes.Complex128DType
 COMPLEX_DTYPE = Literal["complex64", "complex128"]
 COMPLEX = np.complex64 | np.complex128
 STRING_DTYPE = Literal["string"]
 STRING = np.str_
-# BYTES_DTYPE = np.dtypes.BytesDType
 BYTES_DTYPE = Literal["bytes"]
 BYTES = np.bytes_
 
@@ -565,7 +551,7 @@ def to_numpy_shortname(self) -> str:
         }
         return data_type_to_numpy[self]
 
-    def to_numpy(self) -> np.dtype[np.generic]:
+    def to_numpy(self) -> np.dtypes.StringDType | np.dtypes.ObjectDType | np.dtype[np.generic]:
         # note: it is not possible to round trip DataType <-> np.dtype
         # due to the fact that DataType.string and DataType.bytes both
         # generally return np.dtype("O") from this function, even though

src/zarr/strings.py

@@ -1,36 +1,87 @@
-from typing import Any
+"""This module contains utilities for working with string arrays across
+different versions of Numpy.
+"""
+
+from typing import Any, cast
 from warnings import warn
 
 import numpy as np
 
+# STRING_DTYPE is the in-memory datatype that will be used for V3 string arrays
+# when reading data back from Zarr.
+# Any valid string-like datatype should be fine for *setting* data.
+
+STRING_DTYPE: np.dtypes.StringDType | np.dtypes.ObjectDType
+NUMPY_SUPPORTS_VLEN_STRING: bool
+
+
+def cast_array(
+    data: np.ndarray[Any, np.dtype[Any]],
+) -> np.ndarray[Any, np.dtypes.StringDType | np.dtypes.ObjectDType]:
+    raise NotImplementedError
+
+
 try:
-    STRING_DTYPE = np.dtype("T")
+    # this new vlen string dtype was added in NumPy 2.0
+    STRING_DTYPE = np.dtypes.StringDType()
     NUMPY_SUPPORTS_VLEN_STRING = True
-except TypeError:
-    STRING_DTYPE = np.dtype("object")
+
+    def cast_array(
+        data: np.ndarray[Any, np.dtype[Any]],
+    ) -> np.ndarray[Any, np.dtypes.StringDType | np.dtypes.ObjectDType]:
+        out = data.astype(STRING_DTYPE, copy=False)
+        return cast(np.ndarray[Any, np.dtypes.StringDType], out)
+
+except AttributeError:
+    # if not available, we fall back on an object array of strings, as in Zarr < 3
+    STRING_DTYPE = np.dtypes.ObjectDType()
     NUMPY_SUPPORTS_VLEN_STRING = False
 
+    def cast_array(
+        data: np.ndarray[Any, np.dtype[Any]],
+    ) -> np.ndarray[Any, np.dtypes.StringDType | np.dtypes.ObjectDType]:
+        out = data.astype(STRING_DTYPE, copy=False)
+        return cast(np.ndarray[Any, np.dtypes.ObjectDType], out)
+
 
 def cast_to_string_dtype(
     data: np.ndarray[Any, np.dtype[Any]], safe: bool = False
-) -> np.ndarray[Any, np.dtype[Any]]:
+) -> np.ndarray[Any, np.dtypes.StringDType | np.dtypes.ObjectDType]:
+    """Take any data and attempt to cast to to our preferred string dtype.
+
+    data :  np.ndarray
+        The data to cast
+
+    safe : bool
+        If True, do not issue a warning if the data is cast from object to string dtype.
+
+    """
     if np.issubdtype(data.dtype, np.str_):
-        return data
+        # legacy fixed-width string type (e.g. "<U10")
+        return cast_array(data)
+        # out = data.astype(STRING_DTYPE, copy=False)
+        # return cast(np.ndarray[Any, np.dtypes.StringDType | np.dtypes.ObjectDType], out)
+    if NUMPY_SUPPORTS_VLEN_STRING:
+        if np.issubdtype(data.dtype, STRING_DTYPE):
+            # already a valid string variable length string dtype
+            return cast_array(data)
     if np.issubdtype(data.dtype, np.object_):
+        # object arrays require more careful handling
         if NUMPY_SUPPORTS_VLEN_STRING:
             try:
                 # cast to variable-length string dtype, fail if object contains non-string data
                 # mypy says "error: Unexpected keyword argument "coerce" for "StringDType"  [call-arg]"
-                return data.astype(np.dtypes.StringDType(coerce=False), copy=False)  # type: ignore[call-arg]
+                # also: Value of type variable "_ScalarType" of "astype" of "ndarray" cannot be "str"  [type-var]
+                out = data.astype(np.dtypes.StringDType(coerce=False), copy=False)  # type: ignore[call-arg,type-var]
+                return cast_array(out)
             except ValueError as e:
                 raise ValueError("Cannot cast object dtype to string dtype") from e
         else:
-            out = data.astype(np.str_)
             if not safe:
                 warn(
-                    f"Casted object dtype to string dtype {out.dtype}. To avoid this warning, "
+                    "Treating object array as valid string array. To avoid this warning, "
                     "cast the data to a string dtype before passing to Zarr or upgrade to NumPy >= 2.",
                     stacklevel=2,
                 )
-            return out
+            return cast_array(data)
     raise ValueError(f"Cannot cast dtype {data.dtype} to string dtype")

tests/test_strings.py

@@ -0,0 +1,35 @@
+"""Tests for the strings module."""
+
+import numpy as np
+import pytest
+
+from zarr.strings import NUMPY_SUPPORTS_VLEN_STRING, STRING_DTYPE, cast_to_string_dtype
+
+
+def test_string_defaults() -> None:
+    if NUMPY_SUPPORTS_VLEN_STRING:
+        assert STRING_DTYPE == np.dtypes.StringDType()
+    else:
+        assert STRING_DTYPE == np.dtypes.ObjectDType()
+
+
+def test_cast_to_string_dtype() -> None:
+    d1 = np.array(["a", "b", "c"])
+    assert d1.dtype == np.dtype("<U1")
+    d1s = cast_to_string_dtype(d1)
+    assert d1s.dtype == STRING_DTYPE
+
+    with pytest.raises(ValueError, match="Cannot cast dtype |S1"):
+        cast_to_string_dtype(d1.astype("|S1"))
+
+    if NUMPY_SUPPORTS_VLEN_STRING:
+        assert cast_to_string_dtype(d1.astype("T")).dtype == STRING_DTYPE
+        assert cast_to_string_dtype(d1.astype("O")).dtype == STRING_DTYPE
+        with pytest.raises(ValueError, match="Cannot cast object dtype to string dtype"):
+            cast_to_string_dtype(np.array([1, "b", "c"], dtype="O"))
+    else:
+        with pytest.warns():
+            assert cast_to_string_dtype(d1.astype("O")).dtype == STRING_DTYPE
+        with pytest.warns():
+            assert cast_to_string_dtype(np.array([1, "b", "c"], dtype="O")).dtype == STRING_DTYPE
+        assert cast_to_string_dtype(d1.astype("O"), safe=True).dtype == STRING_DTYPE