From 6165f07b7d5fd54113068cccc37303bf2400d80b Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Fri, 5 Jul 2024 16:15:43 +0200
Subject: [PATCH] (fix): move to `_compat` for `ImportError` stuff

---
 src/anndata/experimental/backed/_compat.py    | 42 +++++++++++++++++++
 src/anndata/experimental/backed/_io.py        |  6 +--
 .../experimental/backed/_lazy_arrays.py       | 31 +-------------
 src/anndata/experimental/backed/_xarray.py    |  6 +--
 4 files changed, 48 insertions(+), 37 deletions(-)
 create mode 100644 src/anndata/experimental/backed/_compat.py

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