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

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

@@ -324,8 +324,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
 
@@ -764,19 +763,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

@@ -33,7 +33,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
@@ -121,8 +125,19 @@ def var_names(self) -> pd.Index[str]:
     def obs_names(self) -> pd.Index[str]:
         return self._adata.obs_names
 
-    def __getitem__(self, index: Index) -> Raw:
-        oidx, vidx = self._normalize_indices(index)
+    def __getitem__(self, index: Index | tuple[AnnData, Index]) -> Raw:
+        from .anndata import AnnData
+
+        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):
@@ -133,7 +148,7 @@ def __getitem__(self, index: Index) -> Raw:
         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() == [