add missing module

Author: rabernat

src/zarr/core/strings.py

@@ -0,0 +1,87 @@
+"""This module contains utilities for working with string arrays across
+different versions of Numpy.
+"""
+
+from typing import Any, Union, 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: Union["np.dtypes.StringDType", "np.dtypes.ObjectDType"]
+NUMPY_SUPPORTS_VLEN_STRING: bool
+
+
+def cast_array(
+    data: np.ndarray[Any, np.dtype[Any]],
+) -> np.ndarray[Any, Union["np.dtypes.StringDType", "np.dtypes.ObjectDType"]]:
+    raise NotImplementedError
+
+
+try:
+    # this new vlen string dtype was added in NumPy 2.0
+    STRING_DTYPE = np.dtypes.StringDType()
+    NUMPY_SUPPORTS_VLEN_STRING = True
+
+    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, Union["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, Union["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_):
+        # 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]"
+                # 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:
+            if not safe:
+                warn(
+                    "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 cast_array(data)
+    raise ValueError(f"Cannot cast dtype {data.dtype} to string dtype")