fix(cohort): #20 follow-ups + clear mypy errors

src/sddb/_catalog_schema.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from typing import Any
+
 import pandas as pd
 from pandas.api import types as pdt
 
@@ -149,7 +151,7 @@ class SchemaError(ValueError):
     """Raised when a catalog/gene-index DataFrame violates the schema contract."""
 
 
-def _is_null(v) -> bool:
+def _is_null(v: Any) -> bool:
     """A single cell is null — scalar NaN/NA/None. A list/array cell is never null."""
     return v is None or (not hasattr(v, "__iter__") and pd.isna(v))
 

src/sddb/catalog.py

@@ -297,7 +297,7 @@ def cohort(self, *, uids: list[str] | None = None, study_ids: list[str] | None =
         """
         if (uids is None) == (study_ids is None):
             raise ValueError("pass exactly one of uids= or study_ids=")
-        col, want = ("uid", uids) if uids is not None else ("study_id", study_ids)
+        col, want = ("uid", uids) if uids is not None else ("study_id", study_ids or [])
         if col not in self._df.columns:
             raise ValueError(f"catalog has no {col!r} column")
         want = [str(x) for x in want]

src/sddb/cohort.py

@@ -155,8 +155,10 @@ def check_updates(self, catalog: Catalog, *, warn: bool = True) -> dict[str, str
         return superseded
 
     def _derive(self, df: pd.DataFrame, filter: dict[str, Any]) -> SpatialDataCohort:
-        """A new cohort over ``df`` carrying this cohort's source (used by subsetting/splits)."""
-        return SpatialDataCohort(df, matched=self.matched, source=self._source, filter=filter)
+        """A new cohort over ``df`` carrying this cohort's source and split manifest (used by subsetting/splits)."""
+        sub = SpatialDataCohort(df, matched=self.matched, source=self._source, filter=filter)
+        sub.split_manifest = self.split_manifest
+        return sub
 
     def split(
         self, *, by: str = "study_id", train: float, val: float, test: float, seed: int
@@ -201,13 +203,20 @@ def where(self, mask: pd.Series) -> SpatialDataCohort:
         """Return a new cohort keeping rows where the boolean ``mask`` (aligned to ``members``) is True."""
         return self._derive(self._df[mask.fillna(False).astype(bool)], self._filter)
 
-    def __or__(self, other: SpatialDataCohort) -> SpatialDataCohort:
+    def __or__(self, other: object) -> SpatialDataCohort:
         """Union, deduped by uid (first occurrence wins). Provenance filter is dropped (no longer a pure facet set)."""
+        if not isinstance(other, SpatialDataCohort):
+            return NotImplemented
         df = pd.concat([self._df, other._df]).drop_duplicates("uid", keep="first")
-        return self._derive(df, {})
+        out = self._derive(df, {})
+        if other.split_manifest is not self.split_manifest:
+            out.split_manifest = None  # mixes splits (or split + non-split): no single manifest applies
+        return out
 
-    def __sub__(self, other: SpatialDataCohort) -> SpatialDataCohort:
+    def __sub__(self, other: object) -> SpatialDataCohort:
         """Difference: drop any member whose uid is in ``other``."""
+        if not isinstance(other, SpatialDataCohort):
+            return NotImplemented
         return self._derive(self._df[~self._df["uid"].isin(other._df["uid"])], self._filter)
 
     def citations(self, path: str | Path, *, bib_url: str | None = None) -> Path:
@@ -296,7 +305,7 @@ def __getitem__(self, i: int) -> Dataset: ...
     def __getitem__(self, i: slice) -> SpatialDataCohort: ...
     def __getitem__(self, i: int | slice) -> Dataset | SpatialDataCohort:
         if isinstance(i, slice):
-            return SpatialDataCohort(self._df.iloc[i], matched=self.matched, source=self._source, filter=self._filter)
+            return self._derive(self._df.iloc[i], self._filter)
         return Dataset(self._df.iloc[i])
 
     def __iter__(self) -> Iterator[Dataset]:

src/sddb/manifest.py

@@ -310,7 +310,7 @@ def download(
         warnings.warn(
             f"member {uid}: pinned store missing; FALLING BACK to latest {latest[0]} "
             f"({latest[1]}) - bytes DIFFER from the pinned version",
-            stacklevel=2,
+            stacklevel=3,  # user -> SpatialDataCohort.download -> here
         )
         pairs.append((uid, latest[1]))
     if missing:

src/sddb/mcp/tools.py

@@ -99,7 +99,7 @@ def genes_tool(
     """
     if isinstance(symbols, str):
         symbols = [symbols]
-    facets = {k: v for k, v in (("organism", organism), ("tissue", tissue), ("disease", disease), ("assay", assay)) if v is not None}
+    facets: dict[str, Any] = {k: v for k, v in (("organism", organism), ("tissue", tissue), ("disease", disease), ("assay", assay)) if v is not None}
     cohort = Catalog(catalog_url).query(
         expressing=symbols, mode=mode, validation=None if include_bronze else "pass", **facets
     )

tests/test_cohort_split.py

@@ -33,3 +33,16 @@ def test_fold_carries_split_manifest(tmp_path):
     sm = folds["train"].split_manifest
     assert isinstance(sm, SplitManifest)
     assert sm.write(tmp_path / "split.json").exists()
+
+
+def test_subset_of_fold_keeps_split_manifest():
+    fold = _coh(6).split(by="study_id", train=0.5, val=0.25, test=0.25, seed=0)["train"]
+    sm = fold.split_manifest
+    assert fold[:1].split_manifest is sm
+    assert fold.where(fold.members["uid"] != "").split_manifest is sm
+
+
+def test_union_keeps_manifest_only_within_one_split():
+    folds = _coh(6).split(by="study_id", train=0.5, val=0.25, test=0.25, seed=0)
+    assert (folds["train"] | folds["val"]).split_manifest is folds["train"].split_manifest
+    assert (folds["train"] | _coh(2)).split_manifest is None

tests/test_cohort_subset.py

@@ -55,3 +55,11 @@ def test_empty_filter_is_empty_cohort_not_error():
 def test_slice_preserves_filter():
     coh = SpatialDataCohort(pd.DataFrame({"uid": ["a", "b", "c"], "zarr_url": ["x"] * 3}), filter={"tissue": "breast"})
     assert coh[:2]._filter == {"tissue": "breast"}
+
+
+def test_set_ops_with_non_cohort_raise_typeerror():
+    coh = _coh(["a"], ["breast"])
+    with pytest.raises(TypeError):
+        coh | {"a"}
+    with pytest.raises(TypeError):
+        coh - ["a"]

tests/test_download_fallback.py

@@ -38,6 +38,7 @@ def _latest_for(source, stem_arg):
         coh.download(tmp_path)
     assert captured["pairs"] == [(stem + "0000", "s3://b/latest.zarr")]
     assert any("fall" in str(x.message).lower() or "latest" in str(x.message).lower() for x in w)
+    assert w[0].filename == __file__  # points at the caller, not sddb internals
 
 
 def test_missing_store_no_fallback_raises_nonstrict(monkeypatch, tmp_path):