diff --git a/src/sddb/catalog.py b/src/sddb/catalog.py index 72d103b..21be451 100644 --- a/src/sddb/catalog.py +++ b/src/sddb/catalog.py @@ -165,7 +165,8 @@ def query( license_set: bool | None = None, noncommercial: bool | None = None, search: str | None = None, - expressing: str | None = None, + expressing: str | list[str] | None = None, + mode: str = "all", min_fraction: float | None = None, **facets: str | float | list[str] | tuple[str, ...] | set[str], ) -> SpatialDataCohort: @@ -186,8 +187,10 @@ def query( still imposes its own ``validation_status == "pass"`` filter unless ``validation_status`` was itself matched, so failed rows do not survive a ``search``. expressing - Gene symbol; keep datasets where it is expressed (see ``genes.ranked``). The first use - downloads the gene index (~142 MB, cached). + Gene symbol or Ensembl id, or a list of them; keep datasets where they are expressed (see + ``genes.where_expressed``). The first use downloads the gene index (~142 MB, cached). + mode + With several ``expressing`` genes: ``"all"`` (default, AND) or ``"any"`` (union). min_fraction Minimum fraction of expressing cells for ``expressing``; requires ``expressing``. **facets @@ -202,8 +205,10 @@ def query( If a keyword is not a facet or range column, a facet column is absent from the loaded catalog, a range value is not numeric, or ``min_fraction`` is given without ``expressing``. """ + if mode not in ("all", "any"): + raise ValueError(f"mode must be 'all' or 'any', got {mode!r}") if min_fraction is not None and expressing is None: - raise ValueError("min_fraction requires expressing=") + raise ValueError("min_fraction requires expressing=") equality, ranges = _split_facets(facets) bad = [k for k in equality if k not in schema.FACETS] if bad: @@ -241,14 +246,15 @@ def query( matched = hits.matched or None keep &= df["uid"].isin(hits.to_df()["uid"]) if expressing is not None and keep.any(): # skip the ~142 MB gene index when nothing is left to filter - ranked = self.genes.ranked(expressing, min_fraction=min_fraction, validation=validation) - keep &= df["uid"].isin(ranked["uid"]) + member = self.genes.where_expressed(expressing, mode=mode, min_fraction=min_fraction, validation=None) + keep &= df["uid"].isin(set(member.to_df()["uid"])) params: dict[str, Any] = { "validation": validation, "license_set": license_set, "noncommercial": noncommercial, "search": search, "expressing": expressing, + "mode": mode if expressing is not None else None, "min_fraction": min_fraction, } flt = {k: v for k, v in params.items() if v is not None} diff --git a/src/sddb/cli.py b/src/sddb/cli.py index 8e43113..b701e55 100644 --- a/src/sddb/cli.py +++ b/src/sddb/cli.py @@ -53,10 +53,11 @@ def query( technology: str | None = None, tier: str | None = None, validation: Annotated[str, typer.Option(help="'pass' (default) or 'all'.")] = "pass", - expressing: Annotated[ - str | None, typer.Option(help="Keep datasets expressing this gene symbol (downloads the gene index once).") + gene: Annotated[ + list[str] | None, + typer.Option("--gene", "--expressing", help="Gene symbol or Ensembl id; repeat to require all (downloads the gene index once)."), ] = None, - min_fraction: Annotated[float | None, typer.Option(help="With --expressing: min fraction_obs_detected.")] = None, + min_fraction: Annotated[float | None, typer.Option(help="With --gene: min fraction_obs_detected.")] = None, search: Annotated[str | None, typer.Option(help="Free-text search over facet values.")] = None, min_obs: Annotated[int | None, typer.Option(help="Keep datasets with n_obs >= this.")] = None, min_features: Annotated[int | None, typer.Option(help="Keep datasets with n_features >= this.")] = None, @@ -70,7 +71,8 @@ def query( try: cat = Catalog(catalog) res = cat.query( - validation=val, search=search, expressing=expressing, min_fraction=min_fraction, **facets, **ranges + validation=val, search=search, expressing=gene or None, mode="all", min_fraction=min_fraction, + **facets, **ranges ) except Exception as err: raise _fail(err) from err @@ -160,20 +162,29 @@ def viewer_url(uid: str, catalog: _CatalogOpt = None) -> None: @app.command() def genes( - symbol: Annotated[str, typer.Argument(help="Gene symbol (case-insensitive).")], + genes: Annotated[list[str], typer.Argument(help="Gene symbol(s) or Ensembl id(s), case-insensitive.")], limit: Annotated[int, typer.Option(help="Max datasets to list.")] = 20, + validation: Annotated[str, typer.Option(help="'all' (default, includes bronze) or 'pass'.")] = "all", catalog: _CatalogOpt = None, ) -> None: - """List datasets expressing SYMBOL, ranked by fraction of observations detected. + """List datasets expressing all GENES (one gene: ranked by fraction of observations detected). The first call downloads the cross-dataset gene index (~105 MB, one-time; cached afterwards). """ + val = _validation(validation) try: - df = Catalog(catalog).genes.ranked(symbol).head(limit) + cat = Catalog(catalog) + flags = [c for c in ("validation_status", "tier") if c in cat.to_df().columns] + if len(genes) == 1: + df = cat.genes.ranked(genes[0], validation=val).head(limit) + df = df.merge(cat.to_df()[["uid", *flags]], on="uid", how="left") + else: + df = cat.genes.where_expressed(genes, mode="all", validation=val).to_df().head(limit) + df = df[[c for c in (*_COLS, *flags) if c in df.columns]] except Exception as err: raise _fail(RuntimeError(f"could not load the gene index: {err}")) from err if df.empty: - typer.echo(f"no datasets express {symbol!r}", err=True) + typer.echo(f"no datasets express {' + '.join(genes)!r}", err=True) return _print_table(list(df.columns), [[_fmt(v) for v in r] for r in df.itertuples(index=False)]) diff --git a/src/sddb/genes.py b/src/sddb/genes.py index cb23a95..e0442e7 100644 --- a/src/sddb/genes.py +++ b/src/sddb/genes.py @@ -41,22 +41,59 @@ def __init__(self, catalog: Catalog, url: str | None = None) -> None: except schema.SchemaError as err: raise schema.SchemaError(f"gene index at {self.url} is invalid: {err}") from err self._df = df + # Casefold the token columns once (the index is ~11M rows) so resolve() is a cheap per-token mask. + self._sym = df["symbol"].str.casefold() + self._fid = df["feature_id"].str.casefold() # may hold pd.NA (legacy rows): compare vectorised only + + def resolve(self, tokens: str | list[str], *, min_fraction: float | None = None) -> dict[str, set[str]]: + """Map each token to the uids of datasets expressing it. + + A token (casefolded, stripped) matches a gene-index row by symbol OR Ensembl ``feature_id``; + an ambiguous symbol therefore unions all its feature ids. Unmatched tokens map to an empty set. + Keys keep the original token spelling. + """ + df = self._df + base = ( + (df["fraction_obs_detected"] >= min_fraction).fillna(False).astype(bool) + if min_fraction is not None + else None + ) + sym, fid = self._sym, self._fid + out: dict[str, set[str]] = {} + for tok in [tokens] if isinstance(tokens, str) else tokens: + key = tok.strip().casefold() + mask = ((sym == key) | (fid == key)).fillna(False).astype(bool) + if base is not None: + mask &= base + out[tok] = {str(u) for u in df.loc[mask, "uid"].dropna()} + return out def where_expressed( - self, symbol: str, *, min_fraction: float | None = None, validation: str | None = "pass" + self, + genes: str | list[str], + *, + mode: str = "all", + min_fraction: float | None = None, + validation: str | None = None, ) -> SpatialDataCohort: - """Catalog datasets expressing ``symbol`` (case-insensitive), as a SpatialDataCohort. + """Catalog datasets expressing ``genes`` (symbols or Ensembl ids, case-insensitive), as a cohort. Parameters ---------- - symbol - Gene symbol. + genes + One gene or a list; each may be a symbol or an Ensembl feature id. + mode + ``"all"`` (default) keeps datasets expressing every gene (AND); ``"any"`` the union. min_fraction - Keep only datasets where ``fraction_obs_detected >= min_fraction``. + Per gene, keep only datasets where ``fraction_obs_detected >= min_fraction``. validation - Keep datasets with this ``validation_status`` (default ``"pass"``, as in ``Catalog.query``); ``None`` disables. + Keep datasets with this ``validation_status``; the default ``None`` includes bronze/non-pass + datasets (flagged via the cohort's ``validation_status``/``tier`` columns). """ - uids = set(self.ranked(symbol, min_fraction=min_fraction, validation=validation)["uid"]) + if mode not in ("all", "any"): + raise ValueError(f"mode must be 'all' or 'any', got {mode!r}") + sets = list(self.resolve(genes, min_fraction=min_fraction).values()) + uids = (set.intersection(*sets) if mode == "all" else set.union(*sets)) if sets else set() cat = self._catalog_df(validation) return SpatialDataCohort(cat[cat["uid"].isin(uids)], source=self._catalog._source()) @@ -67,15 +104,16 @@ def _catalog_df(self, validation: str | None) -> pd.DataFrame: return cat if validation is None else cat[cat["validation_status"] == validation] def ranked( - self, symbol: str, *, min_fraction: float | None = None, validation: str | None = "pass" + self, symbol: str, *, min_fraction: float | None = None, validation: str | None = None ) -> pd.DataFrame: """Per-dataset ``uid``, ``fraction_obs_detected``, ``total_counts`` for ``symbol``, best detection first. The per-dataset detection profile for ``symbol``. Filters as in :meth:`where_expressed`; one row per - catalog dataset (max over matching features). + catalog dataset (max over matching features). ``symbol`` may be a symbol or an Ensembl feature id. """ df = self._df - mask = df["symbol"].str.casefold() == symbol.casefold() + key = symbol.strip().casefold() + mask = ((self._sym == key) | (self._fid == key)).fillna(False).astype(bool) if min_fraction is not None: mask &= df["fraction_obs_detected"] >= min_fraction cols = ["uid", "fraction_obs_detected", "total_counts"] diff --git a/src/sddb/mcp/server.py b/src/sddb/mcp/server.py index 0eb9e85..1643f14 100644 --- a/src/sddb/mcp/server.py +++ b/src/sddb/mcp/server.py @@ -55,8 +55,19 @@ def query( def describe(uid: str) -> dict[str, Any]: return tools.describe_tool(uid) - def genes(symbol: str) -> list[dict[str, Any]]: - return tools.genes_tool(symbol) + def genes( + symbols: list[str], + organism: str | None = None, + tissue: str | None = None, + disease: str | None = None, + assay: str | None = None, + mode: str = "all", + include_bronze: bool = True, + ) -> list[dict[str, Any]]: + return tools.genes_tool( + symbols, mode=mode, include_bronze=include_bronze, + **_drop_none(organism=organism, tissue=tissue, disease=disease, assay=assay), + ) def facets(field: str | None = None) -> list[str]: return tools.facets_tool(field) @@ -109,7 +120,13 @@ def download_cohort( description="Filter datasets by facets, ranges (min_obs/min_features), gene (expressing) and free-text search.", )(query) server.tool(name="describe", description="Catalog row and element shapes for one dataset uid.")(describe) - server.tool(name="genes", description="Datasets whose gene index contains a gene symbol.")(genes) + server.tool( + name="genes", + description=( + "Datasets expressing one or more genes (symbols or Ensembl ids; mode all=AND, any=union), " + "optionally filtered by organism/tissue/disease/assay. Bronze datasets are included by default." + ), + )(genes) server.tool(name="facets", description="List facet columns, or the distinct values of a facet field.")(facets) server.tool(name="cite", description="BibTeX for the studies of a filtered cohort.")(cite) server.tool(name="download", description="Plan (sizes) or fetch a filtered cohort to the server download dir.")(download_cohort) diff --git a/src/sddb/mcp/tools.py b/src/sddb/mcp/tools.py index bd0c3c0..286c2ed 100644 --- a/src/sddb/mcp/tools.py +++ b/src/sddb/mcp/tools.py @@ -82,9 +82,28 @@ def describe_tool(uid: str, catalog_url: str | None = None) -> dict[str, Any]: return row -def genes_tool(symbol: str, catalog_url: str | None = None) -> list[dict[str, Any]]: - """Return catalog rows of datasets whose gene index contains ``symbol``.""" - return _records(Catalog(catalog_url).genes.datasets_with(symbol).to_df()) +def genes_tool( + symbols: list[str], + organism: str | None = None, + tissue: str | None = None, + disease: str | None = None, + assay: str | None = None, + mode: str = "all", + include_bronze: bool = True, + catalog_url: str | None = None, +) -> list[dict[str, Any]]: + """Catalog rows of datasets expressing ``symbols`` (symbols or Ensembl ids; ``mode`` all/any). + + Bronze/non-pass datasets are included by default (``include_bronze=False`` keeps only validated ones); + ``organism``/``tissue``/``disease``/``assay`` narrow the result. + """ + 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} + cohort = Catalog(catalog_url).query( + expressing=symbols, mode=mode, validation=None if include_bronze else "pass", **facets + ) + return _records(cohort.to_df()) def facets_tool(field: str | None = None, catalog_url: str | None = None) -> list[str]: diff --git a/tests/test_catalog.py b/tests/test_catalog.py index 64fef5f..01ada85 100644 --- a/tests/test_catalog.py +++ b/tests/test_catalog.py @@ -251,6 +251,26 @@ def test_query_expressing(tmp_path): assert sorted(d.uid for d in cat.query(organism="human", expressing="EPCAM")) == ["uid0001", "uid0002", "uid0005"] +def test_query_expressing_multi_gene(tmp_path): + cat = _cat_genes(tmp_path) + assert [d.uid for d in cat.query(expressing=["EPCAM", "KRT8"])] == ["uid0001"] + assert sorted(d.uid for d in cat.query(expressing=["KRT8", "PTPRC"], mode="any")) == ["uid0001", "uid0002"] + assert [d.uid for d in cat.query(expressing=["ENSG00000119888", "krt8"])] == ["uid0001"] + assert cat.query(expressing=["EPCAM", "KRT8"], mode="any")._filter["mode"] == "any" + + +def test_query_expressing_validation_mask_not_doubled(tmp_path): + cat = _cat_genes(tmp_path) + assert len(cat.query(expressing=["Alb"])) == 0 # default validation="pass" drops the failing row + assert [d.uid for d in cat.query(expressing=["Alb"], validation=None)] == ["uid0004"] + + +def test_query_bad_mode_errors(tmp_path): + cat = _cat_genes(tmp_path) + with pytest.raises(ValueError, match="mode must be"): + cat.query(mode="anyy") + + def test_query_expressing_min_fraction(tmp_path): cat = _cat_genes(tmp_path) assert [d.uid for d in cat.query(expressing="EPCAM", min_fraction=0.35)] == ["uid0001"] diff --git a/tests/test_cli.py b/tests/test_cli.py index 2310162..8c0a066 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -98,6 +98,40 @@ def test_query_expressing_and_range_flags(tmp_path): assert "uid0002" not in r.output # excluded by search=lung +def test_query_gene_repeatable_ands(tmp_path): + cat = write_fixture_catalog(tmp_path / "catalog.parquet") + write_fixture_gene_index(tmp_path / "gene_index.parquet") + r = runner.invoke(app, ["query", "--catalog", cat.as_uri(), "--gene", "EPCAM", "--gene", "KRT8", "--json"]) + assert r.exit_code == 0, r.output + assert {d["uid"] for d in json.loads(r.output)} == {"uid0001"} + + +def test_genes_multi_gene_and_validation(tmp_path): + cat = write_fixture_catalog(tmp_path / "catalog.parquet") + write_fixture_gene_index(tmp_path / "gene_index.parquet") + r = runner.invoke(app, ["genes", "EPCAM", "PTPRC", "--catalog", cat.as_uri()]) + assert r.exit_code == 0, r.output + assert "uid0002" in r.output and "uid0001" not in r.output + r = runner.invoke(app, ["genes", "Alb", "--catalog", cat.as_uri()]) + assert "uid0004" in r.output # bronze included by default + r = runner.invoke(app, ["genes", "Alb", "--validation", "pass", "--catalog", cat.as_uri()]) + assert "uid0004" not in r.output + + +def test_genes_ensembl_id_matches_symbol_and_flags_bronze(tmp_path): + cat = write_fixture_catalog(tmp_path / "catalog.parquet") + write_fixture_gene_index(tmp_path / "gene_index.parquet") + by_sym = runner.invoke(app, ["genes", "EPCAM", "--catalog", cat.as_uri()]) + by_id = runner.invoke(app, ["genes", "ENSG00000119888", "--catalog", cat.as_uri()]) + assert by_id.exit_code == 0, by_id.output + assert by_id.output == by_sym.output and "uid0001" in by_id.output + assert "validation_status" in by_id.output + bronze = runner.invoke(app, ["genes", "Alb", "--catalog", cat.as_uri()]) + assert "fail" in bronze.output + multi = runner.invoke(app, ["genes", "EPCAM", "KRT8", "--catalog", cat.as_uri()]) + assert "validation_status" in multi.output and "tier" in multi.output + + def test_query_min_obs_excludes_smaller(cat_url): r = runner.invoke(app, ["query", "--catalog", cat_url, "--organism", "human"]) assert "uid0005" in r.output # n_obs=4200, present without the flag diff --git a/tests/test_cli_genes.py b/tests/test_cli_genes.py index 5e058ed..36932fa 100644 --- a/tests/test_cli_genes.py +++ b/tests/test_cli_genes.py @@ -12,7 +12,7 @@ class _FakeIndex: - def ranked(self, symbol): + def ranked(self, symbol, validation=None): return pd.DataFrame( { "uid": ["uid0002", "uid0001", "uid0003"], diff --git a/tests/test_genes.py b/tests/test_genes.py index dfdc598..2b2386c 100644 --- a/tests/test_genes.py +++ b/tests/test_genes.py @@ -30,8 +30,8 @@ def test_datasets_with_validation(tmp_path): extra["uid"] = "uid0003" pd.concat([gi, extra], ignore_index=True).to_parquet(tmp_path / "gene_index.parquet") cat = Catalog(cat_p.as_uri(), cache_dir=tmp_path / "cache") - assert [d.uid for d in cat.genes.datasets_with("Alb")] == ["uid0003"] - assert sorted(d.uid for d in cat.genes.datasets_with("Alb", validation=None)) == ["uid0003", "uid0004"] + assert [d.uid for d in cat.genes.datasets_with("Alb", validation="pass")] == ["uid0003"] + assert sorted(d.uid for d in cat.genes.datasets_with("Alb")) == ["uid0003", "uid0004"] def test_refresh_passed_to_sidecar(tmp_path, monkeypatch): @@ -58,3 +58,137 @@ def test_where_expressed_is_datasets_with(tmp_path): b = cat.genes.datasets_with("EPCAM").to_df() assert list(a["uid"]) == list(b["uid"]) assert [d.uid for d in cat.genes.where_expressed("EPCAM", min_fraction=0.3)] == ["uid0001", "uid0002"] + + +def _stub_genes(tmp_path): + """Gene index with nullable dtypes: bronze uid0004 expresses EPCAM; MALAT1 is legacy (NA feature_id).""" + import pandas as pd + + cat_p = write_fixture_catalog(tmp_path / "catalog.parquet") + rows = [ + ("EPCAM", "ENSG1", "uid0001", 5000.0, 0.4), + ("EPCAM", "ENSG1", "uid0003", 3000.0, 0.3), + ("EPCAM", "ENSG1", "uid0004", 100.0, 0.1), # uid0004 is validation=fail + ("CD3D", "ENSG2", "uid0001", 900.0, 0.5), + ("MALAT1", None, "uid0002", 50.0, 0.9), + ] + df = pd.DataFrame(rows, columns=["symbol", "feature_id", "uid", "total_counts", "fraction_obs_detected"]) + df = df.astype( + {"symbol": "string", "feature_id": "string", "uid": "string", "total_counts": "Float64", + "fraction_obs_detected": "Float64"} + ) + assert df["feature_id"].isna().sum() == 1 + df.to_parquet(tmp_path / "gene_index.parquet") + return Catalog(cat_p.as_uri(), cache_dir=tmp_path / "cache").genes + + +def _uids(c): + return sorted(d.uid for d in c) + + +def test_symbol_and_ensembl_resolve_same(tmp_path): + g = _stub_genes(tmp_path) + assert _uids(g.where_expressed("EPCAM")) == _uids(g.where_expressed(" ensg1 ")) == ["uid0001", "uid0003", "uid0004"] + assert g.resolve(["epcam", "nope"]) == {"epcam": {"uid0001", "uid0003", "uid0004"}, "nope": set()} + + +def test_and_vs_any(tmp_path): + g = _stub_genes(tmp_path) + assert _uids(g.where_expressed(["EPCAM", "CD3D"], mode="all")) == ["uid0001"] + assert _uids(g.where_expressed(["EPCAM", "CD3D"], mode="any")) == ["uid0001", "uid0003", "uid0004"] + assert _uids(g.where_expressed(["EPCAM", "NOPE"], mode="all")) == [] + assert _uids(g.where_expressed(["EPCAM"], mode="all", min_fraction=0.3)) == ["uid0001", "uid0003"] + + +def test_bronze_default_and_pass_filter(tmp_path): + g = _stub_genes(tmp_path) + df = g.where_expressed("EPCAM").to_df().set_index("uid") + assert "uid0004" in df.index + assert df.loc["uid0004", "validation_status"] == "fail" + assert "tier" in df.columns + assert _uids(g.where_expressed("EPCAM", validation="pass")) == ["uid0001", "uid0003"] + + +def test_legacy_symbol_only_no_crash(tmp_path): + g = _stub_genes(tmp_path) + assert _uids(g.where_expressed("MALAT1")) == ["uid0002"] + assert _uids(g.where_expressed(["MALAT1", "EPCAM"], mode="any")) == ["uid0001", "uid0002", "uid0003", "uid0004"] + + +def test_backcompat_single_shape(tmp_path): + g = _stub_genes(tmp_path) + a, b = g.where_expressed("EPCAM").to_df(), g.datasets_with("EPCAM").to_df() + assert list(a.columns) == list(b.columns) and list(a["uid"]) == list(b["uid"]) + + +# --------------------------------------------------------------------------- +# Cross-surface consistency guard (spec §3, shared canonical rules). +# +# The fixture below and the expected memberships in CANONICAL_MEMBERSHIP are the +# SHARED expectation that sub-project A's engine tests also encode +# (build_gene_entities / build_gene_search in the spatialdata-db repo). The web +# consumes A's compact artifact and the client queries gene_index.parquet directly; +# both must agree on gene -> dataset membership. This is a client-side guard: it +# does NOT import engine code. A change to the canonical rules (spec §3: Ensembl-keyed +# entities, symbol-keyed legacy entities, AND = intersection, bronze included by +# default and flagged, validation="pass" restricts) MUST break BOTH repos' tests. +# +# Shared fixture (nullable "string" dtypes so pd.NA is a real missing feature_id): +# EPCAM / ENSG1 in {uid0001, uid0003} Ensembl-keyed, several datasets +# CD3D / ENSG2 in {uid0001, uid0004} uid0004 is BRONZE (validation "fail") +# MALAT1 / in {uid0002} legacy symbol-only -> sym:malat1 entity +# --------------------------------------------------------------------------- + +CANONICAL_MEMBERSHIP = { + # symbol and Ensembl id resolve to the same entity -> same uid set + "EPCAM": {"uid0001", "uid0003"}, + "ENSG1": {"uid0001", "uid0003"}, + # multi-gene AND = intersection; OR = union + "AND(EPCAM,CD3D)": {"uid0001"}, + "OR(EPCAM,CD3D)": {"uid0001", "uid0003", "uid0004"}, + # bronze: present by default (flagged), excluded under validation="pass" + "CD3D": {"uid0001", "uid0004"}, + "CD3D@pass": {"uid0001"}, + # legacy symbol-only gene (NA feature_id) resolves by its symbol + "MALAT1": {"uid0002"}, +} + + +def _canonical_gene_index(): + import pandas as pd + + rows = [ + ("EPCAM", "ENSG1", "uid0001", 5000.0, 0.4), + ("EPCAM", "ENSG1", "uid0003", 3000.0, 0.3), + ("CD3D", "ENSG2", "uid0001", 900.0, 0.5), + ("CD3D", "ENSG2", "uid0004", 120.0, 0.2), # bronze: uid0004 validation_status "fail" + ("MALAT1", pd.NA, "uid0002", 50.0, 0.9), # legacy symbol-only row + ] + df = pd.DataFrame(rows, columns=["symbol", "feature_id", "uid", "total_counts", "fraction_obs_detected"]) + return df.astype( + {"symbol": "string", "feature_id": "string", "uid": "string", "total_counts": "Float64", + "fraction_obs_detected": "Float64"} + ) + + +def test_cross_surface_canonical_membership(tmp_path): + """Client membership equals the canonical expectation shared with A's engine tests (spec §3).""" + cat_p = write_fixture_catalog(tmp_path / "catalog.parquet") + df = _canonical_gene_index() + assert df["feature_id"].isna().sum() == 1 # legacy row is a real pd.NA + df.to_parquet(tmp_path / "gene_index.parquet") + g = Catalog(cat_p.as_uri(), cache_dir=tmp_path / "cache").genes + + # symbol and Ensembl id -> same entity + assert _uids(g.where_expressed("EPCAM")) == sorted(CANONICAL_MEMBERSHIP["EPCAM"]) + assert _uids(g.where_expressed("ENSG1")) == sorted(CANONICAL_MEMBERSHIP["ENSG1"]) + # multi-gene AND / OR + assert _uids(g.where_expressed(["EPCAM", "CD3D"], mode="all")) == sorted(CANONICAL_MEMBERSHIP["AND(EPCAM,CD3D)"]) + assert _uids(g.where_expressed(["EPCAM", "CD3D"], mode="any")) == sorted(CANONICAL_MEMBERSHIP["OR(EPCAM,CD3D)"]) + # bronze: default includes and flags; validation="pass" excludes + assert _uids(g.where_expressed("CD3D")) == sorted(CANONICAL_MEMBERSHIP["CD3D"]) + cd3d = g.where_expressed("CD3D").to_df().set_index("uid") + assert cd3d.loc["uid0004", "validation_status"] == "fail" + assert _uids(g.where_expressed("CD3D", validation="pass")) == sorted(CANONICAL_MEMBERSHIP["CD3D@pass"]) + # legacy symbol-only gene resolves by symbol + assert _uids(g.where_expressed("MALAT1")) == sorted(CANONICAL_MEMBERSHIP["MALAT1"]) diff --git a/tests/test_mcp_tools.py b/tests/test_mcp_tools.py index 4190690..ae6c320 100644 --- a/tests/test_mcp_tools.py +++ b/tests/test_mcp_tools.py @@ -55,10 +55,38 @@ def test_describe_tool_remote_failure_and_missing(cat_url): tools.describe_tool("nope", cat_url) +def _uids(rows): + return {r["uid"] for r in rows} + + def test_genes_tool(cat_url): - rows = tools.genes_tool("epcam", cat_url) - assert {r["uid"] for r in rows} == {"uid0001", "uid0002", "uid0005"} + rows = tools.genes_tool(["epcam"], catalog_url=cat_url) + assert _uids(rows) == {"uid0001", "uid0002", "uid0005"} assert json.loads(json.dumps(rows)) == rows + assert _uids(tools.genes_tool("epcam", catalog_url=cat_url)) == _uids(rows) # bare string tolerated + + +def test_genes_tool_symbol_and_ensembl_agree(cat_url): + assert _uids(tools.genes_tool(["ENSG00000119888"], catalog_url=cat_url)) == _uids( + tools.genes_tool(["EPCAM"], catalog_url=cat_url) + ) + + +def test_genes_tool_modes(cat_url): + assert _uids(tools.genes_tool(["EPCAM", "KRT8"], catalog_url=cat_url)) == {"uid0001"} + assert _uids(tools.genes_tool(["EPCAM", "PTPRC"], mode="all", catalog_url=cat_url)) == {"uid0002"} + assert _uids(tools.genes_tool(["KRT8", "PTPRC"], mode="any", catalog_url=cat_url)) == {"uid0001", "uid0002"} + + +def test_genes_tool_bronze_default(cat_url): + assert _uids(tools.genes_tool(["Alb"], catalog_url=cat_url)) == {"uid0004"} # validation "fail" row + assert tools.genes_tool(["Alb"], include_bronze=False, catalog_url=cat_url) == [] + + +def test_genes_tool_facets_compose(cat_url): + rows = tools.genes_tool(["EPCAM"], tissue="lung", catalog_url=cat_url) + assert _uids(rows) == {"uid0001"} + assert _uids(tools.genes_tool(["EPCAM"], organism="mouse", catalog_url=cat_url)) == set() def test_server_registers_tools(): @@ -197,3 +225,24 @@ def run(self): ... assert seen["organism"] == "human" assert "min_obs" not in seen assert "n_features__gte" not in seen + + +def test_server_genes_forwards(monkeypatch): + import sddb.mcp.server as server + + seen = {} + monkeypatch.setattr(tools, "genes_tool", lambda symbols, **kw: seen.update(symbols=symbols, **kw) or []) + registered = {} + + class FakeServer: + def __init__(self, name): ... + def tool(self, name, description=""): + def deco(fn): + registered[name] = fn + return fn + return deco + + monkeypatch.setattr(server, "_server_class", lambda: FakeServer) + server.build_server() + registered["genes"](["EPCAM", "CD3D"], tissue="lung", mode="any", include_bronze=False) + assert seen == {"symbols": ["EPCAM", "CD3D"], "mode": "any", "include_bronze": False, "tissue": "lung"}