implement __contains__

Author: ilia-kats

src/mudata/_core/mudata.py

@@ -547,6 +547,14 @@ def __getitem__(self, index) -> AnnData | MuData:
         else:
             return MuData(self, as_view=True, index=index)
 
+    def __contains__(self, key) -> bool:
+        if isinstance(key, str):
+            return key in self._mod
+        elif type(key).__module__.startswith("anndata.acc") and type(key).__name__ == "AdRef":
+            return AnnData.__contains__.__get__(self)(key)
+        else:
+            raise TypeError(f"Unexpected key {key!r}.")
+
     @property
     def mod(self) -> Mapping[str, AnnData | MuData]:
         """Dictionary of modalities."""

tests/test_obs_var.py

@@ -153,6 +153,7 @@ def test_names_make_unique(mdata: md.MuData):
     Version(ad.__version__) < Version("0.13dev0"), reason="anndata version too old, no accessor support"
 )
 def test_accessors(mdata: md.MuData):
+    assert ad.acc.A.obs["arange"] in mdata
     assert (mdata[ad.acc.A.obs["arange"]] == mdata.obs["arange"]).all()
     with pytest.raises(KeyError, match="test"):
         mdata[ad.acc.A.var["test"]]

tests/test_update.py

@@ -146,6 +146,9 @@ def test_update_simple(mdata: MuData, axis: Axis):
     for mod in mdata.mod.keys():
         assert mdata.obsmap[mod].dtype.kind == "u"
         assert mdata.varmap[mod].dtype.kind == "u"
+        assert mod in mdata
+    with pytest.raises(TypeError):
+        1 in mdata  # noqa: B015
 
     # names along non-axis are concatenated
     assert mdata.shape[1 - axis] == sum(mod.shape[1 - axis] for mod in mdata.mod.values())