array api-related upstream-dev failures

xarray/core/dtypes.py

@@ -4,8 +4,9 @@
 from typing import Any
 
 import numpy as np
+from pandas.api.types import is_extension_array_dtype
 
-from xarray.core import utils
+from xarray.core import npcompat, utils
 
 # Use as a sentinel value to indicate a dtype appropriate NA value.
 NA = utils.ReprObject("<NA>")
@@ -60,22 +61,22 @@ def maybe_promote(dtype: np.dtype) -> tuple[np.dtype, Any]:
     # N.B. these casting rules should match pandas
     dtype_: np.typing.DTypeLike
     fill_value: Any
-    if np.issubdtype(dtype, np.floating):
+    if isdtype(dtype, "real floating"):
         dtype_ = dtype
         fill_value = np.nan
-    elif np.issubdtype(dtype, np.timedelta64):
+    elif isinstance(dtype, np.dtype) and np.issubdtype(dtype, np.timedelta64):
         # See https://github.com/numpy/numpy/issues/10685
         # np.timedelta64 is a subclass of np.integer
         # Check np.timedelta64 before np.integer
         fill_value = np.timedelta64("NaT")
         dtype_ = dtype
-    elif np.issubdtype(dtype, np.integer):
+    elif isdtype(dtype, "integral"):
         dtype_ = np.float32 if dtype.itemsize <= 2 else np.float64
         fill_value = np.nan
-    elif np.issubdtype(dtype, np.complexfloating):
+    elif isdtype(dtype, "complex floating"):
         dtype_ = dtype
         fill_value = np.nan + np.nan * 1j
-    elif np.issubdtype(dtype, np.datetime64):
+    elif isinstance(dtype, np.dtype) and np.issubdtype(dtype, np.datetime64):
         dtype_ = dtype
         fill_value = np.datetime64("NaT")
     else:
@@ -118,16 +119,16 @@ def get_pos_infinity(dtype, max_for_int=False):
     -------
     fill_value : positive infinity value corresponding to this dtype.
     """
-    if issubclass(dtype.type, np.floating):
+    if isdtype(dtype, "real floating"):
         return np.inf
 
-    if issubclass(dtype.type, np.integer):
+    if isdtype(dtype, "integral"):
         if max_for_int:
             return np.iinfo(dtype).max
         else:
             return np.inf
 
-    if issubclass(dtype.type, np.complexfloating):
+    if isdtype(dtype, "complex floating"):
         return np.inf + 1j * np.inf
 
     return INF
@@ -146,24 +147,66 @@ def get_neg_infinity(dtype, min_for_int=False):
     -------
     fill_value : positive infinity value corresponding to this dtype.
     """
-    if issubclass(dtype.type, np.floating):
+    if isdtype(dtype, "real floating"):
         return -np.inf
 
-    if issubclass(dtype.type, np.integer):
+    if isdtype(dtype, "integral"):
         if min_for_int:
             return np.iinfo(dtype).min
         else:
             return -np.inf
 
-    if issubclass(dtype.type, np.complexfloating):
+    if isdtype(dtype, "complex floating"):
         return -np.inf - 1j * np.inf
 
     return NINF
 
 
-def is_datetime_like(dtype):
+def is_datetime_like(dtype) -> bool:
     """Check if a dtype is a subclass of the numpy datetime types"""
-    return np.issubdtype(dtype, np.datetime64) or np.issubdtype(dtype, np.timedelta64)
+    return _is_numpy_subdtype(dtype, (np.datetime64, np.timedelta64))
+
+
+def is_object(dtype) -> bool:
+    """Check if a dtype is object"""
+    return _is_numpy_subdtype(dtype, object)
+
+
+def is_string(dtype) -> bool:
+    """Check if a dtype is a string dtype"""
+    return _is_numpy_subdtype(dtype, (np.str_, np.character))
+
+
+def _is_numpy_subdtype(dtype, kind) -> bool:
+    if not isinstance(dtype, np.dtype):
+        return False
+
+    kinds = kind if isinstance(kind, tuple) else (kind,)
+    return any(np.issubdtype(dtype, kind) for kind in kinds)
+
+
+def isdtype(dtype, kind: str | tuple[str, ...], xp=None) -> bool:
+    """Compatibility wrapper for isdtype() from the array API standard.
+
+    Unlike xp.isdtype(), kind must be a string.
+    """
+    # TODO(shoyer): remove this wrapper when Xarray requires
+    # numpy>=2 and pandas extensions arrays are implemented in
+    # Xarray via the array API
+    if not isinstance(kind, str) and not (
+        isinstance(kind, tuple) and all(isinstance(k, str) for k in kind)
+    ):
+        raise TypeError(f"kind must be a string or a tuple of strings: {repr(kind)}")
+
+    if isinstance(dtype, np.dtype):
+        return npcompat.isdtype(dtype, kind)
+    elif is_extension_array_dtype(dtype):
+        # we never want to match pandas extension array dtypes
+        return False
+    else:
+        if xp is None:
+            xp = np
+        return xp.isdtype(dtype, kind)
 
 
 def result_type(
@@ -184,12 +227,26 @@ def result_type(
     -------
     numpy.dtype for the result.
     """
-    types = {np.result_type(t).type for t in arrays_and_dtypes}
+    from xarray.core.duck_array_ops import get_array_namespace
+
+    # TODO(shoyer): consider moving this logic into get_array_namespace()
+    # or another helper function.
+    namespaces = {get_array_namespace(t) for t in arrays_and_dtypes}
+    non_numpy = namespaces - {np}
+    if non_numpy:
+        [xp] = non_numpy
+    else:
+        xp = np
+
+    types = {xp.result_type(t) for t in arrays_and_dtypes}
 
-    for left, right in PROMOTE_TO_OBJECT:
-        if any(issubclass(t, left) for t in types) and any(
-            issubclass(t, right) for t in types
-        ):
-            return np.dtype(object)
+    if any(isinstance(t, np.dtype) for t in types):
+        # only check if there's numpy dtypes – the array API does not
+        # define the types we're checking for
+        for left, right in PROMOTE_TO_OBJECT:
+            if any(np.issubdtype(t, left) for t in types) and any(
+                np.issubdtype(t, right) for t in types
+            ):
+                return xp.dtype(object)
 
-    return np.result_type(*arrays_and_dtypes)
+    return xp.result_type(*arrays_and_dtypes)

xarray/core/duck_array_ops.py

@@ -142,17 +142,25 @@ def fail_on_dask_array_input(values, msg=None, func_name=None):
 
 def isnull(data):
     data = asarray(data)
-    scalar_type = data.dtype.type
-    if issubclass(scalar_type, (np.datetime64, np.timedelta64)):
+
+    xp = get_array_namespace(data)
+    scalar_type = data.dtype
+    if dtypes.is_datetime_like(scalar_type):
         # datetime types use NaT for null
         # note: must check timedelta64 before integers, because currently
         # timedelta64 inherits from np.integer
         return isnat(data)
-    elif issubclass(scalar_type, np.inexact):
+    elif dtypes.isdtype(scalar_type, ("real floating", "complex floating"), xp=xp):
         # float types use NaN for null
         xp = get_array_namespace(data)
         return xp.isnan(data)
-    elif issubclass(scalar_type, (np.bool_, np.integer, np.character, np.void)):
+    elif dtypes.isdtype(scalar_type, ("bool", "integral"), xp=xp) or (
+        isinstance(scalar_type, np.dtype)
+        and (
+            np.issubdtype(scalar_type, np.character)
+            or np.issubdtype(scalar_type, np.void)
+        )
+    ):
         # these types cannot represent missing values
         return full_like(data, dtype=bool, fill_value=False)
     else:
@@ -406,13 +414,22 @@ def f(values, axis=None, skipna=None, **kwargs):
         if invariant_0d and axis == ():
             return values
 
-        values = asarray(values)
+        xp = get_array_namespace(values)
+        values = asarray(values, xp=xp)
 
-        if coerce_strings and values.dtype.kind in "SU":
+        if coerce_strings and dtypes.is_string(values.dtype):
             values = astype(values, object)
 
         func = None
-        if skipna or (skipna is None and values.dtype.kind in "cfO"):
+        if skipna or (
+            skipna is None
+            and (
+                dtypes.isdtype(
+                    values.dtype, ("complex floating", "real floating"), xp=xp
+                )
+                or dtypes.is_object(values.dtype)
+            )
+        ):
             nanname = "nan" + name
             func = getattr(nanops, nanname)
         else:
@@ -477,8 +494,8 @@ def _datetime_nanmin(array):
     - numpy nanmin() don't work on datetime64 (all versions at the moment of writing)
     - dask min() does not work on datetime64 (all versions at the moment of writing)
     """
-    assert array.dtype.kind in "mM"
     dtype = array.dtype
+    assert dtypes.is_datetime_like(dtype)
     # (NaT).astype(float) does not produce NaN...
     array = where(pandas_isnull(array), np.nan, array.astype(float))
     array = min(array, skipna=True)
@@ -515,7 +532,7 @@ def datetime_to_numeric(array, offset=None, datetime_unit=None, dtype=float):
     """
     # Set offset to minimum if not given
     if offset is None:
-        if array.dtype.kind in "Mm":
+        if dtypes.is_datetime_like(array.dtype):
             offset = _datetime_nanmin(array)
         else:
             offset = min(array)
@@ -527,7 +544,7 @@ def datetime_to_numeric(array, offset=None, datetime_unit=None, dtype=float):
     # This map_blocks call is for backwards compatibility.
     # dask == 2021.04.1 does not support subtracting object arrays
     # which is required for cftime
-    if is_duck_dask_array(array) and np.issubdtype(array.dtype, object):
+    if is_duck_dask_array(array) and dtypes.is_object(array.dtype):
         array = array.map_blocks(lambda a, b: a - b, offset, meta=array._meta)
     else:
         array = array - offset
@@ -537,11 +554,11 @@ def datetime_to_numeric(array, offset=None, datetime_unit=None, dtype=float):
         array = np.array(array)
 
     # Convert timedelta objects to float by first converting to microseconds.
-    if array.dtype.kind in "O":
+    if dtypes.is_object(array.dtype):
         return py_timedelta_to_float(array, datetime_unit or "ns").astype(dtype)
 
     # Convert np.NaT to np.nan
-    elif array.dtype.kind in "mM":
+    elif dtypes.is_datetime_like(array.dtype):
         # Convert to specified timedelta units.
         if datetime_unit:
             array = array / np.timedelta64(1, datetime_unit)
@@ -641,7 +658,7 @@ def mean(array, axis=None, skipna=None, **kwargs):
     from xarray.core.common import _contains_cftime_datetimes
 
     array = asarray(array)
-    if array.dtype.kind in "Mm":
+    if dtypes.is_datetime_like(array.dtype):
         offset = _datetime_nanmin(array)
 
         # xarray always uses np.datetime64[ns] for np.datetime64 data
@@ -689,7 +706,9 @@ def cumsum(array, axis=None, **kwargs):
 
 def first(values, axis, skipna=None):
     """Return the first non-NA elements in this array along the given axis"""
-    if (skipna or skipna is None) and values.dtype.kind not in "iSU":
+    if (skipna or skipna is None) and not (
+        dtypes.isdtype(values.dtype, "signed integer") or dtypes.is_string(values.dtype)
+    ):
         # only bother for dtypes that can hold NaN
         if is_chunked_array(values):
             return chunked_nanfirst(values, axis)
@@ -700,7 +719,9 @@ def first(values, axis, skipna=None):
 
 def last(values, axis, skipna=None):
     """Return the last non-NA elements in this array along the given axis"""
-    if (skipna or skipna is None) and values.dtype.kind not in "iSU":
+    if (skipna or skipna is None) and not (
+        dtypes.isdtype(values.dtype, "signed integer") or dtypes.is_string(values.dtype)
+    ):
         # only bother for dtypes that can hold NaN
         if is_chunked_array(values):
             return chunked_nanlast(values, axis)

xarray/core/npcompat.py

@@ -28,3 +28,33 @@
 # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+try:
+    # requires numpy>=2.0
+    from numpy import isdtype  # type: ignore[attr-defined,unused-ignore]
+except ImportError:
+    import numpy as np
+
+    dtype_kinds = {
+        "bool": np.bool_,
+        "signed integer": np.signedinteger,
+        "unsigned integer": np.unsignedinteger,
+        "integral": np.integer,
+        "real floating": np.floating,
+        "complex floating": np.complexfloating,
+        "numeric": np.number,
+    }
+
+    def isdtype(dtype, kind):
+        kinds = kind if isinstance(kind, tuple) else (kind,)
+
+        unknown_dtypes = [kind for kind in kinds if kind not in dtype_kinds]
+        if unknown_dtypes:
+            raise ValueError(f"unknown dtype kinds: {unknown_dtypes}")
+
+        # verified the dtypes already, no need to check again
+        translated_kinds = [dtype_kinds[kind] for kind in kinds]
+        if isinstance(dtype, np.generic):
+            return any(isinstance(dtype, kind) for kind in translated_kinds)
+        else:
+            return any(np.issubdtype(dtype, kind) for kind in translated_kinds)