From 44c84bfbd8f20c4f50b816b7a342ed9400c51ddd Mon Sep 17 00:00:00 2001 From: John Kerl Date: Fri, 17 Jan 2025 12:03:45 -0500 Subject: [PATCH] [python] Pass more context in `io/shaping.py` (#3582) --- apis/python/src/tiledbsoma/io/shaping.py | 54 ++++++++++++++++++------ 1 file changed, 42 insertions(+), 12 deletions(-) diff --git a/apis/python/src/tiledbsoma/io/shaping.py b/apis/python/src/tiledbsoma/io/shaping.py index 450ed323f5..777a0bb9c9 100644 --- a/apis/python/src/tiledbsoma/io/shaping.py +++ b/apis/python/src/tiledbsoma/io/shaping.py @@ -98,6 +98,7 @@ def show_experiment_shapes( uri, visitor=_leaf_visitor_show_shapes, args=args, + context=context, ) return ok @@ -149,7 +150,7 @@ def upgrade_experiment_shapes( # they're from the experiment as-is. nobs = None nvars = {} - with tiledbsoma.Experiment.open(uri) as exp: + with tiledbsoma.Experiment.open(uri, context=context) as exp: if "obs" in exp: nobs = exp.obs.count if "ms" in exp: @@ -171,6 +172,7 @@ def upgrade_experiment_shapes( uri, visitor=_leaf_visitor_upgrade, args=args, + context=context, ) return ok @@ -258,7 +260,7 @@ def resize_experiment( # If the user hasn't provided a key -- e.g. a from-anndata-append-with-resize # on one measurement while the experiment's other measurements aren't being # updated -- then we need to find those other measurements' var-shapes. - with tiledbsoma.Experiment.open(uri) as exp: + with tiledbsoma.Experiment.open(uri, context=context) as exp: for ms_key in exp.ms.keys(): if ms_key not in nvars.keys(): nvars[ms_key] = exp.ms[ms_key].var._maybe_soma_joinid_shape or 1 @@ -267,6 +269,7 @@ def resize_experiment( uri, visitor=_leaf_visitor_resize, args=args, + context=context, ) return ok @@ -277,26 +280,39 @@ def _treewalk( node_name: str | None = None, visitor: Any, args: SizingArgs, + context: tiledbsoma.SOMATileDBContext | None, ) -> bool: retval = True - with tiledbsoma.open(uri) as item: + with tiledbsoma.open(uri, context=context) as item: if isinstance(item, tiledbsoma.Experiment): if "obs" in item: ok = _treewalk( - item["obs"].uri, node_name="obs", visitor=visitor, args=args + item["obs"].uri, + node_name="obs", + visitor=visitor, + args=args, + context=context, ) retval = retval and ok if "ms" in item: ok = _treewalk( - item["ms"].uri, node_name="ms", visitor=visitor, args=args + item["ms"].uri, + node_name="ms", + visitor=visitor, + args=args, + context=context, ) retval = retval and ok elif isinstance(item, tiledbsoma.Measurement): if "var" in item: ok = _treewalk( - item["var"].uri, node_name="var", visitor=visitor, args=args + item["var"].uri, + node_name="var", + visitor=visitor, + args=args, + context=context, ) retval = retval and ok @@ -310,6 +326,7 @@ def _treewalk( node_name=coll_name, visitor=visitor, args=args, + context=context, ) retval = retval and ok @@ -320,11 +337,17 @@ def _treewalk( if node_name == "ms": args["ms_name"] = key - ok = _treewalk(item[key].uri, node_name=key, visitor=visitor, args=args) + ok = _treewalk( + item[key].uri, + node_name=key, + visitor=visitor, + args=args, + context=context, + ) retval = retval and ok else: - ok = visitor(item, node_name=node_name, args=args) + ok = visitor(item, node_name=node_name, args=args, context=context) retval = retval and ok return retval @@ -335,6 +358,7 @@ def _leaf_visitor_show_shapes( *, node_name: str, args: SizingArgs, + context: tiledbsoma.SOMATileDBContext | None, ) -> bool: retval = True if isinstance(item, tiledbsoma.DataFrame): @@ -366,6 +390,7 @@ def _leaf_visitor_upgrade( *, node_name: str, args: SizingArgs, + context: tiledbsoma.SOMATileDBContext | None, ) -> bool: verbose = args["verbose"] check_only = args["check_only"] @@ -392,7 +417,7 @@ def _leaf_visitor_upgrade( f" Applying tiledbsoma_upgrade_soma_joinid_shape({count})", file=args["output_handle"], ) - with tiledbsoma.DataFrame.open(item.uri, "w") as writer: + with tiledbsoma.DataFrame.open(item.uri, "w", context=context) as writer: writer.tiledbsoma_upgrade_soma_joinid_shape(count) else: if verbose: @@ -427,7 +452,9 @@ def _leaf_visitor_upgrade( f" Applying tiledbsoma_upgrade_shape({new_shape})", file=args["output_handle"], ) - with tiledbsoma.SparseNDArray.open(item.uri, "w") as writer: + with tiledbsoma.SparseNDArray.open( + item.uri, "w", context=context + ) as writer: writer.tiledbsoma_upgrade_shape(new_shape) else: if verbose: @@ -448,6 +475,7 @@ def _leaf_visitor_resize( *, node_name: str, args: SizingArgs, + context: tiledbsoma.SOMATileDBContext | None, ) -> bool: verbose = args["verbose"] check_only = args["check_only"] @@ -487,7 +515,7 @@ def _leaf_visitor_resize( f" Applying tiledbsoma_resize_soma_joinid_shape({new_soma_joinid_shape})", file=args["output_handle"], ) - with tiledbsoma.DataFrame.open(item.uri, "w") as writer: + with tiledbsoma.DataFrame.open(item.uri, "w", context=context) as writer: writer.tiledbsoma_resize_soma_joinid_shape(new_soma_joinid_shape) elif isinstance(item, tiledbsoma.SparseNDArray): @@ -504,7 +532,9 @@ def _leaf_visitor_resize( else: if verbose: print(f" Applying resize({new_shape})", file=args["output_handle"]) - with tiledbsoma.SparseNDArray.open(item.uri, "w") as writer: + with tiledbsoma.SparseNDArray.open( + item.uri, "w", context=context + ) as writer: writer.resize(new_shape) elif isinstance(item, tiledbsoma.DenseNDArray):