Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Backport release-1.15] [python] Pass more context in io/shaping.py #3587

Merged
merged 1 commit into from
Jan 17, 2025
Merged

[Backport release-1.15] [python] Pass more context in io/shaping.py #3587

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 42 additions & 12 deletions apis/python/src/tiledbsoma/io/shaping.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@ def show_experiment_shapes(
uri,
visitor=_leaf_visitor_show_shapes,
args=args,
context=context,
)
return ok

Expand Down Expand Up @@ -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:
Expand All @@ -171,6 +172,7 @@ def upgrade_experiment_shapes(
uri,
visitor=_leaf_visitor_upgrade,
args=args,
context=context,
)
return ok

Expand Down Expand Up @@ -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
Expand All @@ -267,6 +269,7 @@ def resize_experiment(
uri,
visitor=_leaf_visitor_resize,
args=args,
context=context,
)
return ok

Expand All @@ -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

Expand All @@ -310,6 +326,7 @@ def _treewalk(
node_name=coll_name,
visitor=visitor,
args=args,
context=context,
)
retval = retval and ok

Expand All @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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"]
Expand All @@ -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:
Expand Down Expand Up @@ -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:
Expand All @@ -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"]
Expand Down Expand Up @@ -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):
Expand All @@ -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):
Expand Down