replace is_duck_array with _arrayfunction_or_api instead

Author: andersy005

xarray/core/arithmetic.py

@@ -14,12 +14,9 @@
     VariableOpsMixin,
 )
 from xarray.core.common import ImplementsArrayReduce, ImplementsDatasetReduce
-from xarray.core.ops import (
-    IncludeNumpySameMethods,
-    IncludeReduceMethods,
-)
+from xarray.core.ops import IncludeNumpySameMethods, IncludeReduceMethods
 from xarray.core.options import OPTIONS, _get_keep_attrs
-from xarray.namedarray.pycompat import is_duck_array
+from xarray.namedarray._typing import _arrayfunction_or_api
 
 
 class SupportsArithmetic:
@@ -48,7 +45,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
         # See the docstring example for numpy.lib.mixins.NDArrayOperatorsMixin.
         out = kwargs.get("out", ())
         for x in inputs + out:
-            if not is_duck_array(x) and not isinstance(
+            if not isinstance(x, _arrayfunction_or_api) and not isinstance(
                 x, self._HANDLED_TYPES + (SupportsArithmetic,)
             ):
                 return NotImplemented

xarray/core/dataset.py

@@ -112,17 +112,14 @@
     broadcast_variables,
     calculate_dimensions,
 )
+from xarray.namedarray._typing import _arrayfunction_or_api
 from xarray.namedarray.daskmanager import DaskManager
 from xarray.namedarray.parallelcompat import get_chunked_array_type, guess_chunkmanager
-from xarray.namedarray.pycompat import (
-    array_type,
-    is_chunked_array,
-)
+from xarray.namedarray.pycompat import array_type, is_chunked_array
 from xarray.namedarray.utils import (
     consolidate_dask_from_array_kwargs,
     either_dict_or_kwargs,
     is_dict_like,
-    is_duck_array,
     is_duck_dask_array,
 )
 from xarray.plot.accessor import DatasetPlotAccessor
@@ -2719,7 +2716,7 @@ def _validate_indexers(
             elif isinstance(v, Sequence) and len(v) == 0:
                 yield k, np.empty((0,), dtype="int64")
             else:
-                if not is_duck_array(v):
+                if not isinstance(v, _arrayfunction_or_api):
                     v = np.asarray(v)
 
                 if v.dtype.kind in "US":

xarray/core/duck_array_ops.py

@@ -16,17 +16,11 @@
 import pandas as pd
 from numpy import all as array_all  # noqa
 from numpy import any as array_any  # noqa
-from numpy import (  # noqa
+from numpy import (
     around,  # noqa
-    einsum,
-    gradient,
     isclose,
-    isin,
     isnat,
-    take,
-    tensordot,
-    transpose,
-    unravel_index,
+    take,  # noqa
     zeros_like,  # noqa
 )
 from numpy import concatenate as _concatenate
@@ -35,9 +29,10 @@
 
 from xarray.core import dask_array_ops, dtypes, nputils
 from xarray.core.utils import module_available
+from xarray.namedarray._typing import _arrayfunction_or_api
 from xarray.namedarray.parallelcompat import get_chunked_array_type, is_chunked_array
 from xarray.namedarray.pycompat import array_type
-from xarray.namedarray.utils import is_duck_array, is_duck_dask_array
+from xarray.namedarray.utils import is_duck_dask_array
 
 dask_available = module_available("dask")
 
@@ -193,7 +188,7 @@ def astype(data, dtype, **kwargs):
 
 
 def asarray(data, xp=np):
-    return data if is_duck_array(data) else xp.asarray(data)
+    return data if isinstance(data, _arrayfunction_or_api) else xp.asarray(data)
 
 
 def as_shared_dtype(scalars_or_arrays, xp=np):

xarray/core/formatting.py

@@ -19,8 +19,8 @@
 from xarray.core.duck_array_ops import array_equiv, astype
 from xarray.core.indexing import MemoryCachedArray
 from xarray.core.options import OPTIONS, _get_boolean_with_default
+from xarray.namedarray._typing import _arrayfunction_or_api
 from xarray.namedarray.pycompat import array_type, to_duck_array, to_numpy
-from xarray.namedarray.utils import is_duck_array
 
 if TYPE_CHECKING:
     from xarray.core.coordinates import AbstractCoordinates
@@ -630,7 +630,7 @@ def short_data_repr(array):
     internal_data = getattr(array, "variable", array)._data
     if isinstance(array, np.ndarray):
         return short_array_repr(array)
-    elif is_duck_array(internal_data):
+    elif isinstance(internal_data, _arrayfunction_or_api):
         return limit_lines(repr(array.data), limit=40)
     elif getattr(array, "_in_memory", None):
         return short_array_repr(array)
@@ -789,7 +789,9 @@ def extra_items_repr(extra_keys, mapping, ab_side, kwargs):
             is_variable = True
         except AttributeError:
             # compare attribute value
-            if is_duck_array(a_mapping[k]) or is_duck_array(b_mapping[k]):
+            if isinstance(a_mapping[k], _arrayfunction_or_api) or isinstance(
+                b_mapping[k], _arrayfunction_or_api
+            ):
                 compatible = array_equiv(a_mapping[k], b_mapping[k])
             else:
                 compatible = a_mapping[k] == b_mapping[k]

xarray/core/indexing.py

@@ -24,13 +24,10 @@
     is_scalar,
     to_0d_array,
 )
+from xarray.namedarray._typing import _arrayfunction_or_api
 from xarray.namedarray.parallelcompat import get_chunked_array_type, is_chunked_array
 from xarray.namedarray.pycompat import array_type, integer_types
-from xarray.namedarray.utils import (
-    either_dict_or_kwargs,
-    is_duck_array,
-    is_duck_dask_array,
-)
+from xarray.namedarray.utils import either_dict_or_kwargs, is_duck_dask_array
 
 if TYPE_CHECKING:
     from numpy.typing import DTypeLike
@@ -377,7 +374,7 @@ def __init__(self, key):
                 k = int(k)
             elif isinstance(k, slice):
                 k = as_integer_slice(k)
-            elif is_duck_array(k):
+            elif isinstance(k, _arrayfunction_or_api):
                 if not np.issubdtype(k.dtype, np.integer):
                     raise TypeError(
                         f"invalid indexer array, does not have integer dtype: {k!r}"
@@ -424,7 +421,7 @@ def __init__(self, key):
                     "Please pass a numpy array by calling ``.compute``. "
                     "See https://github.com/dask/dask/issues/8958."
                 )
-            elif is_duck_array(k):
+            elif isinstance(k, _arrayfunction_or_api):
                 if not np.issubdtype(k.dtype, np.integer):
                     raise TypeError(
                         f"invalid indexer array, does not have integer dtype: {k!r}"

xarray/core/nputils.py

@@ -14,7 +14,7 @@
     from numpy import RankWarning
 
 from xarray.core.options import OPTIONS
-from xarray.namedarray.pycompat import is_duck_array
+from xarray.namedarray._typing import _arrayfunction_or_api
 
 try:
     import bottleneck as bn
@@ -143,7 +143,10 @@ def _advanced_indexer_subspaces(key):
 
     non_slices = [k for k in key if not isinstance(k, slice)]
     broadcasted_shape = np.broadcast_shapes(
-        *[item.shape if is_duck_array(item) else (0,) for item in non_slices]
+        *[
+            item.shape if isinstance(item, _arrayfunction_or_api) else (0,)
+            for item in non_slices
+        ]
     )
     ndim = len(broadcasted_shape)
     mixed_positions = advanced_index_positions[0] + np.arange(ndim)

xarray/core/variable.py

@@ -36,17 +36,13 @@
     infix_dims,
     maybe_coerce_to_str,
 )
+from xarray.namedarray._typing import _arrayfunction_or_api
 from xarray.namedarray.core import NamedArray
 from xarray.namedarray.parallelcompat import get_chunked_array_type
-from xarray.namedarray.pycompat import (
-    integer_types,
-    is_0d_dask_array,
-    is_chunked_array,
-)
+from xarray.namedarray.pycompat import integer_types, is_0d_dask_array, is_chunked_array
 from xarray.namedarray.utils import (
     either_dict_or_kwargs,
     is_dict_like,
-    is_duck_array,
     is_duck_dask_array,
 )
 
@@ -411,7 +407,7 @@ def data(self):
         Variable.as_numpy
         Variable.values
         """
-        if is_duck_array(self._data):
+        if isinstance(self._data, _arrayfunction_or_api):
             return self._data
         elif isinstance(self._data, indexing.ExplicitlyIndexed):
             return self._data.get_duck_array()
@@ -636,7 +632,7 @@ def _validate_indexers(self, key):
         for dim, k in zip(self.dims, key):
             if not isinstance(k, BASIC_INDEXING_TYPES):
                 if not isinstance(k, Variable):
-                    if not is_duck_array(k):
+                    if not isinstance(k, _arrayfunction_or_api):
                         k = np.asarray(k)
                     if k.ndim > 1:
                         raise IndexError(
@@ -681,7 +677,7 @@ def _broadcast_indexes_outer(self, key):
             if isinstance(k, Variable):
                 k = k.data
             if not isinstance(k, BASIC_INDEXING_TYPES):
-                if not is_duck_array(k):
+                if not isinstance(k, _arrayfunction_or_api):
                     k = np.asarray(k)
                 if k.size == 0:
                     # Slice by empty list; numpy could not infer the dtype
@@ -940,7 +936,7 @@ def load(self, **kwargs):
             self._data = as_compatible_data(loaded_data)
         elif isinstance(self._data, ExplicitlyIndexed):
             self._data = self._data.get_duck_array()
-        elif not is_duck_array(self._data):
+        elif not isinstance(self._data, _arrayfunction_or_api):
             self._data = np.asarray(self._data)
         return self
 

xarray/namedarray/pycompat.py

@@ -8,7 +8,8 @@
 from packaging.version import Version
 
 from xarray.core.utils import is_scalar
-from xarray.namedarray.utils import is_duck_array, is_duck_dask_array
+from xarray.namedarray._typing import _arrayfunction_or_api
+from xarray.namedarray.utils import is_duck_dask_array
 
 integer_types = (int, np.integer)
 
@@ -85,7 +86,9 @@ def mod_version(mod: ModType) -> Version:
 
 
 def is_chunked_array(x) -> bool:
-    return is_duck_dask_array(x) or (is_duck_array(x) and hasattr(x, "chunks"))
+    return is_duck_dask_array(x) or (
+        isinstance(x, _arrayfunction_or_api) and hasattr(x, "chunks")
+    )
 
 
 def is_0d_dask_array(x):
@@ -120,7 +123,7 @@ def to_duck_array(data):
 
     if isinstance(data, ExplicitlyIndexed):
         return data.get_duck_array()
-    elif is_duck_array(data):
+    elif isinstance(data, _arrayfunction_or_api):
         return data
     else:
         return np.asarray(data)

xarray/namedarray/utils.py

@@ -7,6 +7,8 @@
 
 import numpy as np
 
+from xarray.namedarray._typing import _arrayfunction_or_api
+
 if TYPE_CHECKING:
     if sys.version_info >= (3, 10):
         from typing import TypeGuard
@@ -15,7 +17,7 @@
 
     from numpy.typing import NDArray
 
-    from xarray.namedarray._typing import T_DuckArray, duckarray
+    from xarray.namedarray._typing import duckarray
 
     try:
         from dask.array.core import Array as DaskArray
@@ -26,10 +28,6 @@
 
 T = TypeVar("T")
 
-# https://stackoverflow.com/questions/74633074/how-to-type-hint-a-generic-numpy-array
-T_DType_co = TypeVar("T_DType_co", bound=np.dtype[np.generic], covariant=True)
-T_DType = TypeVar("T_DType", bound=np.dtype[np.generic])
-
 
 # Singleton type, as per https://github.com/python/typing/pull/240
 class Default(Enum):
@@ -68,21 +66,7 @@ def is_dask_collection(x: object) -> TypeGuard[DaskCollection]:
 
 
 def is_duck_dask_array(x: duckarray[Any, Any]) -> TypeGuard[DaskArray]:
-    return is_dask_collection(x)
-
-
-def is_duck_array(value: Any) -> TypeGuard[T_DuckArray]:
-    if isinstance(value, np.ndarray):
-        return True
-    return (
-        hasattr(value, "ndim")
-        and hasattr(value, "shape")
-        and hasattr(value, "dtype")
-        and (
-            (hasattr(value, "__array_function__") and hasattr(value, "__array_ufunc__"))
-            or hasattr(value, "__array_namespace__")
-        )
-    )
+    return isinstance(x, _arrayfunction_or_api) and is_dask_collection(x)
 
 
 def to_0d_object_array(

xarray/testing.py

@@ -13,7 +13,7 @@
 from xarray.core.dataset import Dataset
 from xarray.core.indexes import Index, PandasIndex, PandasMultiIndex, default_indexes
 from xarray.core.variable import IndexVariable, Variable
-from xarray.namedarray.utils import is_duck_array
+from xarray.namedarray._typing import _arrayfunction_or_api
 
 __all__ = (
     "assert_allclose",
@@ -229,14 +229,14 @@ def assert_duckarray_equal(x, y, err_msg="", verbose=True):
     """Like `np.testing.assert_array_equal`, but for duckarrays"""
     __tracebackhide__ = True
 
-    if not is_duck_array(x) and not utils.is_scalar(x):
+    if not isinstance(x, _arrayfunction_or_api) and not utils.is_scalar(x):
         x = np.asarray(x)
 
-    if not is_duck_array(y) and not utils.is_scalar(y):
+    if not isinstance(y, _arrayfunction_or_api) and not utils.is_scalar(y):
         y = np.asarray(y)
 
-    if (is_duck_array(x) and utils.is_scalar(y)) or (
-        utils.is_scalar(x) and is_duck_array(y)
+    if (isinstance(x, _arrayfunction_or_api) and utils.is_scalar(y)) or (
+        utils.is_scalar(x) and isinstance(y, _arrayfunction_or_api)
     ):
         equiv = (x == y).all()
     else: