basic support for AnnData accessors

this should keep scanpy plotting working with MuData objects

Author: ilia-kats

pyproject.toml

@@ -48,7 +48,7 @@ dev = [
   "pre-commit",
   "twine>=4.0.2",
 ]
-test = [ "coverage>=7.10", "mudata[io]", "pytest" ]
+test = [ "coverage>=7.10", "mudata[io]", "packaging", "pytest" ]
 doc = [
   "docutils>=0.8,!=0.18.*,!=0.19.*",
   "ipykernel",

src/mudata/_core/mudata.py

@@ -513,6 +513,25 @@ def strings_to_categoricals(self, df: pd.DataFrame | None = None) -> pd.DataFram
     def __getitem__(self, index) -> AnnData | MuData:
         if isinstance(index, str):
             return self._mod[index]
+        elif type(index).__module__.startswith("anndata.acc") and type(index).__name__ == "AdRef":
+            try:
+                return index.acc.get(self, index.idx)
+            except KeyError as e:
+                if index.acc.dim in ("obs", "var"):
+                    for modname, mod in self._mod.items():
+                        try:
+                            index.acc.get(mod, index.idx)
+                        except KeyError:
+                            pass
+                        else:
+                            raise KeyError(
+                                f"There is no key {index.idx} in MuData .{index.acc.dim} but there is one in {modname} .{index.acc.dim}. Consider running `pull_{index.acc.dim}()` to update global .{index.acc.dim}."
+                            ) from e
+                    raise KeyError(
+                        f"There is no key {index.idx} in MuData .{index.acc.dim} or in .{index.acc.dim} of any modalities."
+                    ) from e
+                else:
+                    raise
         else:
             return MuData(self, as_view=True, index=index)
 

tests/test_obs_var.py

@@ -1,8 +1,10 @@
 from pathlib import Path
 
+import anndata as ad
 import numpy as np
 import pandas as pd
 import pytest
+from packaging.version import Version
 
 import mudata as md
 
@@ -145,3 +147,14 @@ def test_names_make_unique(mdata: md.MuData):
 
     with pytest.raises(TypeError, match="axis="):
         getattr(mdata, f"{attr}_names_make_unique")()
+
+
+@pytest.mark.skipif(
+    Version(ad.__version__) < Version("0.13dev0"), reason="anndata version too old, no accessor support"
+)
+def test_accessors(mdata: md.MuData):
+    assert (mdata[ad.acc.A.obs["arange"]] == mdata.obs["arange"]).all()
+    with pytest.raises(KeyError, match="any modalities"):
+        mdata[ad.acc.A.var["test"]]
+    with pytest.raises(KeyError, match="there is one in"):
+        mdata[ad.acc.A.var["assert-bool"]]