parse all

Author: flying-sheep

src/anndata/acc/__init__.py

@@ -381,7 +381,29 @@ def resolve(self, spec: str, *, strict: Literal[True] = True) -> P: ...
     @overload
     def resolve(self, spec: str, *, strict: Literal[False]) -> P | None: ...
     def resolve(self, spec: str, *, strict: bool = True) -> P | None:
-        """Create accessor from string."""
+        """Create accessor from string.
+
+        Examples
+        --------
+        >>> A.resolve("X[:,:]")
+        A[:, :]
+        >>> A.resolve("layers.y[c,:]")
+        A.layers['y']['c', :]
+        >>> A.resolve("layers.y[:,g]")
+        A.layers['y'][:, 'g']
+        >>> A.resolve("obs.a")
+        A.obs['a']
+        >>> A.resolve("var.b")
+        A.var['b']
+        >>> A.resolve("obsm.c.0")
+        A.obsm['c'][:, 0]
+        >>> A.resolve("varm.d.1")
+        A.varm['d'][:, 1]
+        >>> A.resolve("obsp.g[c1,:]")
+        A.obsp['g']['c1', :]
+        >>> A.resolve("obsp.g[:,c2]")
+        A.obsp['g'][:, 'c2']
+        """
         from ._parse import parse
 
         return parse(self, spec, strict=strict)

src/anndata/acc/_parse.py

@@ -6,9 +6,10 @@
 from . import GraphAcc, LayerAcc, MetaVecAcc, MultiAcc
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
     from typing import Literal
 
-    from . import AdAcc, AdPath, Idx2D
+    from . import AdAcc, AdPath, GraphVecAcc, Idx2D, LayerVecAcc
 
 
 @overload
@@ -23,21 +24,19 @@ def parse[P: AdPath](a: AdAcc[P], spec: str, *, strict: bool = True) -> P | None
         except ValueError:
             return None
 
+    if spec.startswith("X["):
+        return _parse_path_2d(lambda _: a, spec)
     if "." not in spec:
         msg = f"Cannot parse accessor {spec!r}"
         raise ValueError(msg)
-    acc, rest = spec.split(".", 1)
-    match getattr(a, acc, None):
-        # TODO: X
-        case LayerAcc() as layers:
-            return _parse_path_layer(layers, rest)
+    acc_name, rest = spec.split(".", 1)
+    match getattr(a, acc_name, None):
+        case LayerAcc() | GraphAcc() as acc:
+            return _parse_path_2d(acc.__getitem__, rest)
         case MetaVecAcc() as meta:
             return meta[rest]
         case MultiAcc() as multi:
             return _parse_path_multi(multi, rest)
-        case GraphAcc():
-            msg = "TODO"
-            raise NotImplementedError(msg)
         case None:  # pragma: no cover
             msg = (
                 f"Unknown accessor {spec!r}. "
@@ -48,14 +47,17 @@ def parse[P: AdPath](a: AdAcc[P], spec: str, *, strict: bool = True) -> P | None
     raise AssertionError(msg)  # pragma: no cover
 
 
-def _parse_path_layer[P: AdPath](layers: LayerAcc[P], spec: str) -> P:
+def _parse_path_2d[P: AdPath](
+    get_vec_acc: Callable[[str], LayerVecAcc[P] | GraphVecAcc[P]], spec: str
+) -> P:
     if not (
         m := re.fullmatch(r"([^\[]+)\[([^,]+),\s?([^\]]+)\]", spec)
     ):  # pragma: no cover
-        msg = f"Cannot parse layer accessor {spec!r}: should be `name[i,:]`/`name[:,j]`"
+        msg = f"Cannot parse accessor {spec!r}: should be `name[i,:]`/`name[:,j]`"
         raise ValueError(msg)
-    layer, i, j = m.groups("")  # "" just for typing
-    return layers[layer][_parse_idx_2d(i, j, str)]
+    name, i, j = m.groups("")  # "" just for typing
+    vec_acc = get_vec_acc(name)
+    return vec_acc[_parse_idx_2d(i, j, str)]
 
 
 def _parse_path_multi[P: AdPath](multi: MultiAcc[P], spec: str) -> P:
@@ -68,6 +70,8 @@ def _parse_path_multi[P: AdPath](multi: MultiAcc[P], spec: str) -> P:
 
 def _parse_idx_2d[Idx: int | str](i: str, j: str, cls: type[Idx]) -> Idx2D[Idx]:
     match i, j:
+        case ":", ":":
+            return slice(None), slice(None)
         case _, ":":
             return cls(i), slice(None)
         case ":", _:

tests/test_accessors.py

@@ -107,29 +107,6 @@ def test_axes(ad_path: AdPath, axes: Collection[Literal["obs", "var"]]) -> None:
     assert ad_path.axes == axes
 
 
-@pytest.mark.parametrize(
-    ("spec", "expected"),
-    [
-        # TODO: pytest.param("???", A[:, :], id="x"),
-        pytest.param("layers.y[c,:]", A.layers["y"]["c", :], id="layer-obs"),
-        pytest.param("layers.y[:,g]", A.layers["y"][:, "g"], id="layer-var"),
-        pytest.param("obs.a", A.obs["a"], id="obs"),
-        pytest.param("var.b", A.var["b"], id="var"),
-        pytest.param("obsm.c.0", A.obsm["c"][:, 0], id="obsm"),
-        pytest.param("varm.d.1", A.varm["d"][:, 1], id="varm"),
-        pytest.param("obsp.g[c1,:]", A.obsp["d"]["c1", :], id="obsp"),
-        pytest.param("obsp.g[:,c2]", A.obsp["d"][:, "c2"], id="varp"),
-    ],
-)
-def test_resolve(spec: str, expected: AdPath) -> None:
-    try:
-        assert A.resolve(spec) == expected
-        assert spec == expected
-        assert expected != ""
-    except NotImplementedError:
-        pytest.xfail("not implemented")
-
-
 @pytest.mark.parametrize(
     "mk_path",
     [