Support for __array_function__ implementers (sparse arrays) [WIP]

xarray/core/duck_array_ops.py

@@ -76,6 +76,11 @@ def isnull(data):
         return _isnull(data)
     except TypeError:
         return np.zeros(data.shape, dtype=bool)
+    except Exception:
+        if hasattr(data, '__array_function__'):
+            return np.isnan(data)
+        else:
+            raise
 
 
 transpose = _dask_or_eager_func('transpose')
@@ -121,7 +126,11 @@ def trapz(y, x, axis):
 
 
 def asarray(data):
-    return data if isinstance(data, dask_array_type) else np.asarray(data)
+    return (
+        data if (isinstance(data, dask_array_type)
+                 or hasattr(data, '__array_function__'))
+        else np.asarray(data)
+    )
 
 
 def as_shared_dtype(scalars_or_arrays):

xarray/core/formatting.py

@@ -357,7 +357,10 @@ def set_numpy_options(*args, **kwargs):
 
 
 def short_array_repr(array):
-    array = np.asarray(array)
+
+    if not hasattr(array, '__array_function__'):
+        array = np.asarray(array)
+
     # default to lower precision so a full (abbreviated) line can fit on
     # one line with the default display_width
     options = {

xarray/core/indexing.py

@@ -656,6 +656,9 @@ def as_indexable(array):
         return PandasIndexAdapter(array)
     if isinstance(array, dask_array_type):
         return DaskIndexingAdapter(array)
+    if hasattr(array, '__array_function__'):
+        return NdArrayLikeIndexAdapter(array)
+
     raise TypeError('Invalid array type: {}'.format(type(array)))
 
 
@@ -1180,6 +1183,15 @@ def __setitem__(self, key, value):
         array[key] = value
 
 
+class NdArrayLikeIndexAdapter(NumpyIndexingAdapter):
+    def __init__(self, array):
+        if not hasattr(array, '__array_function__'):
+            raise TypeError(
+                'NdArrayLikeIndexAdapter must wrap an object that implements '
+                'the __array_function__ protocol')
+        self.array = array
+
+
 class DaskIndexingAdapter(ExplicitlyIndexedNDArrayMixin):
     """Wrap a dask array to support explicit indexing."""
 

xarray/core/variable.py

@@ -179,6 +179,9 @@ def as_compatible_data(data, fastpath=False):
         else:
             data = np.asarray(data)
 
+    if hasattr(data, '__array_function__'):
+        return data
+
     # validate whether the data is valid data types
     data = np.asarray(data)
 
@@ -207,6 +210,8 @@ def _as_array_or_item(data):
 
     TODO: remove this (replace with np.asarray) once these issues are fixed
     """
+    if hasattr(data, '__array_function__'):
+        return data
     data = np.asarray(data)
     if data.ndim == 0:
         if data.dtype.kind == 'M':