Class LazilyIndexedArray

xarray/core/indexing.py:707–812

Wrap an array to make basic and outer indexing lazy.

class LazilyIndexedArray(ExplicitlyIndexedNDArrayMixin):
    """Wrap an array to make basic and outer indexing lazy."""

    __slots__ = ("_shape", "array", "key")

    def __init__(self, array: Any, key: ExplicitIndexer | None = None):
        """
        Parameters
        ----------
        array : array_like
            Array like object to index.
        key : ExplicitIndexer, optional
            Array indexer. If provided, it is assumed to already be in
            canonical expanded form.
        """
        if isinstance(array, type(self)) and key is None:
            # unwrap
            key = array.key  # type: ignore[has-type, unused-ignore]
            array = array.array  # type: ignore[has-type, unused-ignore]

        if key is None:
            key = BasicIndexer((slice(None),) * array.ndim)

        self.array = as_indexable(array)
        self.key = key

        shape: _Shape = ()
        for size, k in zip(self.array.shape, self.key.tuple, strict=True):
            if isinstance(k, slice):
                shape += (len(range(*k.indices(size))),)
            elif isinstance(k, np.ndarray):
                shape += (k.size,)
        self._shape = shape

    def _updated_key(self, new_key: ExplicitIndexer) -> BasicIndexer | OuterIndexer:
        iter_new_key = iter(expanded_indexer(new_key.tuple, self.ndim))

        full_key: list[OuterIndexerType] = []
        for size, k in zip(self.array.shape, self.key.tuple, strict=True):
            if isinstance(k, integer_types):
                full_key.append(k)
            else:
                full_key.append(_index_indexer_1d(k, next(iter_new_key), size))
        full_key_tuple = tuple(full_key)

        if all(isinstance(k, integer_types + (slice,)) for k in full_key_tuple):
            return BasicIndexer(cast(tuple[BasicIndexerType, ...], full_key_tuple))
        return OuterIndexer(full_key_tuple)

    @property
    def shape(self) -> _Shape:
        return self._shape

    def get_duck_array(self):
        from xarray.backends.common import BackendArray

        if isinstance(self.array, BackendArray):
            array = self.array[self.key]