Avoid calling np.asarray on lazy indexing classes

xarray/coding/variables.py

@@ -9,7 +9,7 @@
 import pandas as pd
 
 from ..core import dtypes, duck_array_ops, indexing
-from ..core.pycompat import is_duck_dask_array
+from ..core.pycompat import is_duck_array, is_duck_dask_array
 from ..core.variable import Variable
 
 
@@ -69,7 +69,10 @@ def __getitem__(self, key):
  return type(self)(self.array[key], self.func, self.dtype)
 
  def __array__(self, dtype=None):
- return self.func(self.array)
+ return self.get_duck_array()
+
+ def get_duck_array(self):
+ return self.func(self.array.get_duck_array())
 
  def __repr__(self):
  return "{}({!r}, func={!r}, dtype={!r})".format(
@@ -216,7 +219,10 @@ def decode(self, variable, name=None):
 
 
 def _scale_offset_decoding(data, scale_factor, add_offset, dtype):
- data = np.array(data, dtype=dtype, copy=True)
+ if not is_duck_array(data):
+ data = np.array(data, dtype=dtype, copy=True)
+ else:
+ data = data.astype(dtype=dtype, copy=True)
  if scale_factor is not None:
  data *= scale_factor
  if add_offset is not None:

xarray/core/common.py

@@ -155,6 +155,9 @@ def __complex__(self: Any) -> complex:
  def __array__(self: Any, dtype: DTypeLike = None) -> np.ndarray:
  return np.asarray(self.values, dtype=dtype)
 
+ def get_duck_array(self):
+ return self.data
+
  def __repr__(self) -> str:
  return formatting.array_repr(self)
 

xarray/core/formatting.py

@@ -15,7 +15,7 @@
 from pandas.errors import OutOfBoundsDatetime
 
 from .duck_array_ops import array_equiv
-from .indexing import MemoryCachedArray
+from .indexing import ExplicitlyIndexed, MemoryCachedArray
 from .options import OPTIONS, _get_boolean_with_default
 from .pycompat import dask_array_type, sparse_array_type
 from .utils import is_duck_array
@@ -516,8 +516,13 @@ def limit_lines(string: str, *, limit: int):
  return string
 
 
-def short_numpy_repr(array):
- array = np.asarray(array)
+def short_array_repr(array):
+ from .common import AbstractArray
+
+ if isinstance(array, (ExplicitlyIndexed, AbstractArray)):
+ array = array.get_duck_array()
+ if not is_duck_array(array):
+ array = np.asarray(array)
 
  # default to lower precision so a full (abbreviated) line can fit on
  # one line with the default display_width
@@ -541,11 +546,11 @@ def short_data_repr(array):
  """Format "data" for DataArray and Variable."""
  internal_data = getattr(array, "variable", array)._data
  if isinstance(array, np.ndarray):
- return short_numpy_repr(array)
+ return short_array_repr(array)
  elif is_duck_array(internal_data):
  return limit_lines(repr(array.data), limit=40)
  elif array._in_memory or array.size < 1e5:
- return short_numpy_repr(array)
+ return short_array_repr(array)
  else:
  # internal xarray array type
  return f"[{array.size} values with dtype={array.dtype}]"
@@ -765,7 +770,7 @@ def diff_array_repr(a, b, compat):
  equiv = array_equiv
 
  if not equiv(a.data, b.data):
- temp = [wrap_indent(short_numpy_repr(obj), start=" ") for obj in (a, b)]
+ temp = [wrap_indent(short_array_repr(obj), start=" ") for obj in (a, b)]
  diff_data_repr = [
  ab_side + "\n" + ab_data_repr
  for ab_side, ab_data_repr in zip(("L", "R"), temp)

xarray/core/indexing.py

@@ -18,7 +18,13 @@
 from .npcompat import DTypeLike
 from .nputils import NumpyVIndexAdapter
 from .options import OPTIONS
-from .pycompat import dask_version, integer_types, is_duck_dask_array, sparse_array_type
+from .pycompat import (
+ dask_version,
+ integer_types,
+ is_duck_array,
+ is_duck_dask_array,
+ sparse_array_type,
+)
 from .types import T_Xarray
 from .utils import (
  NDArrayMixin,
@@ -442,8 +448,11 @@ class ExplicitlyIndexedNDArrayMixin(NDArrayMixin, ExplicitlyIndexed):
  __slots__ = ()
 
  def __array__(self, dtype=None):
+ return np.asarray(self.get_duck_array(), dtype=dtype)
+
+ def get_duck_array(self):
  key = BasicIndexer((slice(None),) * self.ndim)
- return np.asarray(self[key], dtype=dtype)
+ return self[key]
 
 
 class ImplicitToExplicitIndexingAdapter(NDArrayMixin):
@@ -458,6 +467,9 @@ def __init__(self, array, indexer_cls=BasicIndexer):
  def __array__(self, dtype=None):
  return np.asarray(self.array, dtype=dtype)
 
+ def get_duck_array(self):
+ return self.array.get_duck_array()
+
  def __getitem__(self, key):
  key = expanded_indexer(key, self.ndim)
  result = self.array[self.indexer_cls(key)]
@@ -520,8 +532,17 @@ def shape(self) -> tuple[int, ...]:
  return tuple(shape)
 
  def __array__(self, dtype=None):
+ return np.asarray(self.get_duck_array(), dtype=dtype)
+
+ def get_duck_array(self):
  array = as_indexable(self.array)
- return np.asarray(array[self.key], dtype=None)
+ array = array[self.key]
+ if isinstance(array, ExplicitlyIndexed):
+ array = array.get_duck_array()
+ if not is_duck_array(array):
+ # This hack is necessary for 0D netCDF4 variables
+ array = np.asarray(array)
+ return array
 
  def transpose(self, order):
  return LazilyVectorizedIndexedArray(self.array, self.key).transpose(order)
@@ -622,6 +643,9 @@ def _ensure_copied(self):
  def __array__(self, dtype=None):
  return np.asarray(self.array, dtype=dtype)
 
+ def get_duck_array(self):
+ return self.array.get_duck_array()
+
  def __getitem__(self, key):
  return type(self)(_wrap_numpy_scalars(self.array[key]))
 
@@ -653,6 +677,9 @@ def __array__(self, dtype=None):
  self._ensure_cached()
  return np.asarray(self.array, dtype=dtype)
 
+ def get_duck_array(self):
+ return self.array.get_duck_array()
+
  def __getitem__(self, key):
  return type(self)(_wrap_numpy_scalars(self.array[key]))
 
@@ -1436,6 +1463,9 @@ def __array__(self, dtype: DTypeLike = None) -> np.ndarray:
  array = array.astype("object")
  return np.asarray(array.values, dtype=dtype)
 
+ def get_duck_array(self) -> np.ndarray:
+ return np.asarray(self)
+
  @property
  def shape(self) -> tuple[int, ...]:
  return (len(self.array),)
@@ -1572,9 +1602,9 @@ def _repr_inline_(self, max_width: int) -> str:
  return format_array_flat(self._get_array_subset(), max_width)
 
  def _repr_html_(self) -> str:
- from .formatting import short_numpy_repr
+ from .formatting import short_array_repr
 
- array_repr = short_numpy_repr(self._get_array_subset())
+ array_repr = short_array_repr(self._get_array_subset())
  return f"<pre>{escape(array_repr)}</pre>"
 
  def copy(self, deep: bool = True) -> PandasMultiIndexingAdapter:

xarray/core/variable.py

@@ -356,6 +356,8 @@ def _in_memory(self):
  def data(self):
  if is_duck_array(self._data):
  return self._data
+ elif isinstance(self._data, indexing.ExplicitlyIndexed):
+ return self._data.get_duck_array()
  else:
  return self.values
 
@@ -463,6 +465,8 @@ def load(self, **kwargs):
  """
  if is_duck_dask_array(self._data):
  self._data = as_compatible_data(self._data.compute(**kwargs))
+ elif isinstance(self._data, indexing.ExplicitlyIndexed):
+ self._data = self._data.get_duck_array()
  elif not is_duck_array(self._data):
  self._data = np.asarray(self._data)
  return self

xarray/tests/test_formatting.py

@@ -486,7 +486,7 @@ def test_set_numpy_options() -> None:
  assert np.get_printoptions() == original_options
 
 
-def test_short_numpy_repr() -> None:
+def test_short_array_repr() -> None:
  cases = [
  np.random.randn(500),
  np.random.randn(20, 20),
@@ -496,16 +496,16 @@ def test_short_numpy_repr() -> None:
  ]
  # number of lines:
  # for default numpy repr: 167, 140, 254, 248, 599
- # for short_numpy_repr: 1, 7, 24, 19, 25
+ # for short_array_repr: 1, 7, 24, 19, 25
  for array in cases:
- num_lines = formatting.short_numpy_repr(array).count("\n") + 1
+ num_lines = formatting.short_array_repr(array).count("\n") + 1
  assert num_lines < 30
 
  # threshold option (default: 200)
  array2 = np.arange(100)
- assert "..." not in formatting.short_numpy_repr(array2)
+ assert "..." not in formatting.short_array_repr(array2)
  with xr.set_options(display_values_threshold=10):
- assert "..." in formatting.short_numpy_repr(array2)
+ assert "..." in formatting.short_array_repr(array2)
 
 
 def test_large_array_repr_length() -> None: