Class AccessibleAsDuckArrayDataStore

xarray/tests/test_dataset.py:254–277

Store that returns a duck array, not convertible to numpy array, on read. Modeled after nVIDIA's kvikio.

class AccessibleAsDuckArrayDataStore(backends.InMemoryDataStore):
    """
    Store that returns a duck array, not convertible to numpy array,
    on read. Modeled after nVIDIA's kvikio.
    """

    def __init__(self):
        super().__init__()
        self._indexvars = set()

    def store(self, variables, *args, **kwargs) -> None:
        super().store(variables, *args, **kwargs)
        for k, v in variables.items():
            if isinstance(v, IndexVariable):
                self._indexvars.add(k)

    def get_variables(self) -> dict[Any, xr.Variable]:
        def lazy_accessible(k, v) -> xr.Variable:
            if k in self._indexvars:
                return v
            data = indexing.LazilyIndexedArray(DuckBackendArrayWrapper(v.values))
            return Variable(v.dims, data, v.attrs)

        return {k: lazy_accessible(k, v) for k, v in self._variables.items()}


class TestDataset: