fix: check .raw length on setting (#2351)

Author: flying-sheep

docs/release-notes/2351.fix.md

@@ -0,0 +1 @@
+Check `Raw` length on creation and fix associated `.adata` in `Raw` slicing {user}`P Angerer`

src/anndata/_core/anndata.py

@@ -340,8 +340,7 @@ def _init_as_view(
         # set raw, easy, as it’s immutable anyways...
         if adata_ref._raw is not None:
             # slicing along variables axis is ignored
-            self._raw = adata_ref.raw[oidx]
-            self._raw._adata = self
+            self._raw = adata_ref.raw[self, oidx]
         else:
             self._raw = None
 
@@ -699,19 +698,20 @@ def raw(self) -> Raw:
         return self._raw
 
     @raw.setter
-    def raw(self, value: AnnData):
+    def raw(self, value: AnnData) -> None:
         if value is None:
             del self.raw
-        elif not isinstance(value, AnnData):
+            return
+        if not isinstance(value, AnnData):
             msg = "Can only init raw attribute with an AnnData object."
             raise ValueError(msg)
-        else:
-            if self.is_view:
-                self._init_as_actual(self.copy())
-            self._raw = Raw(self, X=value.X, var=value.var, varm=value.varm)
+        raw = Raw(self, X=value.X, var=value.var, varm=value.varm)
+        if self.is_view:
+            self._init_as_actual(self.copy())
+        self._raw = raw
 
     @raw.deleter
-    def raw(self):
+    def raw(self) -> None:
         if self.is_view:
             self._init_as_actual(self.copy())
         self._raw = None

src/anndata/_core/raw.py

@@ -35,7 +35,11 @@ def __init__(
         X: np.ndarray | CSMatrix | None = None,
         var: pd.DataFrame | Mapping[str, Sequence] | None = None,
         varm: AxisArrays | Mapping[str, np.ndarray] | None = None,
-    ):
+    ) -> None:
+        if X is not None and X.shape[0] != adata.n_obs:
+            msg = f"X has {X.shape[0]} rows, but n_obs is {adata.n_obs}"
+            raise ValueError(msg)
+
         self._adata = adata
         self._n_obs = adata.n_obs
         # construct manually
@@ -126,14 +130,26 @@ def obs_names(self) -> pd.Index[str]:
     @overload
     def __getitem__(self, index: AdRef) -> InMemoryArray: ...
     @overload
-    def __getitem__(self, index: Index) -> Raw: ...
-    def __getitem__(self, index: Index | AdRef) -> Raw | InMemoryArray:
+    def __getitem__(self, index: Index | tuple[AnnData, Index]) -> Raw: ...
+    def __getitem__(
+        self, index: Index | tuple[AnnData, Index] | AdRef
+    ) -> Raw | InMemoryArray:
         from ..acc import AdRef
+        from .anndata import AnnData
 
         if isinstance(index, AdRef):
             return index.acc.get(self, index.idx)  # type: ignore  # no official Raw support here
 
-        oidx, vidx = self._normalize_indices(index)
+        if (
+            isinstance(index, tuple)
+            and len(index) == 2
+            and isinstance(index[0], AnnData)
+        ):
+            adata, index = index
+            oidx, vidx = self._normalize_indices(index)
+        else:
+            oidx, vidx = self._normalize_indices(index)
+            adata = self._adata[oidx]
 
         # To preserve two dimensional shape
         if isinstance(vidx, int | np.integer):
@@ -144,7 +160,7 @@ def __getitem__(self, index: Index | AdRef) -> Raw | InMemoryArray:
         X = _subset(self.X, (oidx, vidx)) if not self._adata.isbacked else None
 
         var = self._var.iloc[vidx]
-        new = Raw(self._adata, X=X, var=var)
+        new = Raw(adata, X=X, var=var)
         if self.varm is not None:
             # Since there is no view of raws
             new.varm = self.varm._view(_RawViewHack(self, vidx), (vidx,)).copy()

tests/test_raw.py

@@ -71,6 +71,11 @@ def test_raw_set_as_none(adata_raw: ad.AnnData):
     assert_equal(a, b)
 
 
+def test_raw_set_error(adata_raw: ad.AnnData) -> None:
+    with pytest.raises(ValueError, match=r"X has 2 rows, but n_obs is 3"):
+        adata_raw.raw = adata_raw[:2].copy()
+
+
 def test_raw_of_view(adata_raw: ad.AnnData):
     adata_view = adata_raw[adata_raw.obs["oanno1"] == "cat2"]
     assert adata_view.raw.X.tolist() == [