diff --git a/src/anndata/experimental/backed/_compat.py b/src/anndata/experimental/backed/_compat.py new file mode 100644 index 000000000..037306698 --- /dev/null +++ b/src/anndata/experimental/backed/_compat.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +try: + from xarray import DataArray +except ImportError: + + class DataArray: + def __repr__(self) -> str: + return "mock DataArray" + + +try: + import xarray as xr +except ImportError: + xr = None + + +try: + from xarray.backends.zarr import ZarrArrayWrapper +except ImportError: + + class ZarrArrayWrapper: + def __repr__(self) -> str: + return "mock ZarrArrayWrapper" + + +try: + from xarray.backends import BackendArray +except ImportError: + + class BackendArray: + def __repr__(self) -> str: + return "mock BackendArray" + + +try: + from xarray import Dataset +except ImportError: + + class Dataset: + def __repr__(self) -> str: + return "mock Dataset" diff --git a/src/anndata/experimental/backed/_io.py b/src/anndata/experimental/backed/_io.py index c092be88d..2aa752e92 100644 --- a/src/anndata/experimental/backed/_io.py +++ b/src/anndata/experimental/backed/_io.py @@ -11,17 +11,13 @@ import dask.array as da import h5py - -try: - import xarray as xr -except ImportError: - xr = None import zarr from ..._core.anndata import AnnData from ..._core.sparse_dataset import sparse_dataset from ...compat import DaskArray from .. import read_dispatched +from ._compat import xr from ._lazy_arrays import CategoricalArray, MaskedArray from ._xarray import Dataset2D diff --git a/src/anndata/experimental/backed/_lazy_arrays.py b/src/anndata/experimental/backed/_lazy_arrays.py index 99d722115..89df81f01 100644 --- a/src/anndata/experimental/backed/_lazy_arrays.py +++ b/src/anndata/experimental/backed/_lazy_arrays.py @@ -5,39 +5,12 @@ import pandas as pd -try: - import xarray as xr - from xarray import DataArray -except ImportError: - xr = None - - class DataArray: - def __repr__(self) -> str: - return "mock DataArray" - - -try: - from xarray.backends.zarr import ZarrArrayWrapper -except ImportError: - - class ZarrArrayWrapper: - def __repr__(self) -> str: - return "mock ZarrArrayWrapper" - - -try: - from xarray.backends import BackendArray -except ImportError: - - class BackendArray: - def __repr__(self) -> str: - return "mock BackendArray" - - from anndata._core.index import Index, _subset from anndata._core.views import as_view from anndata.compat import H5Array, ZarrArray +from ._compat import BackendArray, DataArray, ZarrArrayWrapper, xr + K = TypeVar("K", bound=Union[H5Array, ZarrArray]) diff --git a/src/anndata/experimental/backed/_xarray.py b/src/anndata/experimental/backed/_xarray.py index 59b04b7a0..6d3b0665f 100644 --- a/src/anndata/experimental/backed/_xarray.py +++ b/src/anndata/experimental/backed/_xarray.py @@ -7,12 +7,12 @@ import pandas as pd -import xarray as xr from ..._core.anndata import _gen_dataframe, _remove_unused_categories from ..._core.file_backing import to_memory from ..._core.index import Index, _subset from ..._core.views import as_view +from ._compat import DataArray, Dataset def get_index_dim(ds): @@ -22,7 +22,7 @@ def get_index_dim(ds): return list(ds.sizes.keys())[0] -class Dataset2D(xr.Dataset): +class Dataset2D(Dataset): @property def index(self) -> pd.Index: coord = list(self.coords.keys())[0] @@ -48,7 +48,7 @@ def __getitem__(self, idx): @_subset.register(Dataset2D) -def _(a: xr.DataArray, subset_idx: Index): +def _(a: DataArray, subset_idx: Index): key = get_index_dim(a) if ( isinstance(subset_idx, tuple) and len(subset_idx) == 1