From 0d2648dc086fef6650e4d19d83e6433b5b10c947 Mon Sep 17 00:00:00 2001 From: Ryan Williams Date: Fri, 20 Dec 2024 12:36:04 -0500 Subject: [PATCH] [python] Make `ExperimentAxisQuery` inherit `somacore`, restore its `.close` (#3493) * make EAQ inherit somacore, restore EAQ.close * restore `test_query_cleanup` sans threadpool checks * Update apis/python/src/tiledbsoma/_query.py Co-authored-by: John Kerl --------- Co-authored-by: John Kerl --- apis/python/src/tiledbsoma/_query.py | 7 +++++-- apis/python/tests/test_experiment_query.py | 20 ++++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) 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"])],