diff --git a/apis/python/src/tiledbsoma/_query.py b/apis/python/src/tiledbsoma/_query.py index 8e6849c187..e7270172b0 100644 --- a/apis/python/src/tiledbsoma/_query.py +++ b/apis/python/src/tiledbsoma/_query.py @@ -188,7 +188,7 @@ def to_anndata(self) -> AnnData: ) -class ExperimentAxisQuery: +class ExperimentAxisQuery(query.ExperimentAxisQuery): """Axis-based query against a SOMA Experiment. ExperimentAxisQuery allows easy selection and extraction of data from a @@ -567,7 +567,10 @@ def to_spatialdata( # type: ignore[no-untyped-def] return sdata - # Context management + # Context management. + # Currently a no-op, but, part of public API so we retain it. + def close(self) -> None: + pass def __enter__(self) -> Self: return self diff --git a/apis/python/tests/test_experiment_query.py b/apis/python/tests/test_experiment_query.py index e70605cf2e..ffb3f17aa7 100644 --- a/apis/python/tests/test_experiment_query.py +++ b/apis/python/tests/test_experiment_query.py @@ -545,6 +545,26 @@ def test_error_corners(soma_experiment: Experiment): next(query.varp(lyr_name)) +@pytest.mark.parametrize("n_obs,n_vars", [(1001, 99)]) +def test_query_cleanup(soma_experiment: soma.Experiment): + """ + Verify soma.Experiment.query works as context manager and stand-alone, + and that it cleans up correctly. + """ + from contextlib import closing + + context = SOMATileDBContext() + soma_experiment = get_soma_experiment_with_context(soma_experiment, context) + + with soma_experiment.axis_query("RNA") as query: + assert query.n_obs == 1001 + assert query.n_vars == 99 + assert query.to_anndata("raw") is not None + + with closing(soma_experiment.axis_query("RNA")) as query: + assert query.to_anndata("raw") is not None + + @pytest.mark.parametrize( "n_obs,n_vars,obsp_layer_names,varp_layer_names,obsm_layer_names,varm_layer_names", [(1001, 99, ["foo"], ["bar"], ["baz"], ["quux"])],