From f92c2f6de923fd0acf8e92c548f00d150b68c999 Mon Sep 17 00:00:00 2001 From: Liam Huber Date: Wed, 30 Sep 2026 08:55:43 -0700 Subject: [PATCH] Extend the storage tool and widget to work with all result types. Also allow the starting point inside the h5 file to be specified. Co-authored-by: Claude Signed-off-by: Liam Huber --- src/flowrep/retrospective/storage.py | 91 ++++++++++++++------- src/flowrep/retrospective/storage_widget.py | 33 +++++--- tests/flowrep_static/node_data.py | 67 +++++++++++++++ tests/unit/test_storage.py | 90 +++++++++++++++++++- tests/unit/test_storage_widget.py | 46 ++++++++++- 5 files changed, 281 insertions(+), 46 deletions(-) create mode 100644 tests/flowrep_static/node_data.py diff --git a/src/flowrep/retrospective/storage.py b/src/flowrep/retrospective/storage.py index 3adc16a2..1045e4aa 100644 --- a/src/flowrep/retrospective/storage.py +++ b/src/flowrep/retrospective/storage.py @@ -1,7 +1,8 @@ """ Convenience tools for accessing :cls:`flowrep.datastructures.LiveWorkflow` data stored in *bagofholding* `H5Bag` objects using "lexical" paths (node names, "inputs"/"outputs", -and port names). +and port names). The data may be the bag's top-level object, or be kept somewhere inside +it, e.g. as an attribute of some other saved object. """ from __future__ import annotations @@ -25,6 +26,25 @@ from bagofholding import H5Bag +DEFAULT_STORAGE_ROOT = "object" +"""Where a bag made by saving node data directly keeps that data.""" + +_NODE_DATA_TYPES = ( + datastructures.AtomicData, + datastructures.ConstantData, + datastructures.DagData, + datastructures.ForEachData, + datastructures.IfData, + datastructures.TryData, + datastructures.WhileData, +) +_LOADABLE_TYPES = ( + *_NODE_DATA_TYPES, + datastructures.InputDataPort, + datastructures.OutputDataPort, +) + + class LexicalBagBrowser: """ A convenience class for browsing and loading data from @@ -32,19 +52,26 @@ class LexicalBagBrowser: Lets you access data using the "lexical" paths (i.e. "."-joined paths of node names, "inputs/outputs", and port names) instead of the actual H5 path inside the file. + Node data kept inside some other saved object can be browsed by passing the storage + path to it as `storage_root`. """ @_import_alarm - def __init__(self, bag: H5Bag | str | pathlib.Path): + def __init__( + self, + bag: H5Bag | str | pathlib.Path, + storage_root: str = DEFAULT_STORAGE_ROOT, + ): if isinstance(bag, (str, pathlib.Path)): self.bag = boh.H5Bag(bag) else: self.bag = bag - validate_bag(self.bag) + self.storage_root = _normalize_root(storage_root) + validate_bag(self.bag, self.storage_root) def list_paths(self) -> list[str]: """A list of all available lexical content paths.""" - return list_lexical_paths(self.bag) + return list_lexical_paths(self.bag, self.storage_root) def widget(self) -> storage_widget.LexicalBagTree: """A jupyter-notebook widget for graphical browsing""" @@ -61,6 +88,7 @@ def load( self, path: str ) -> ( datastructures.AtomicData + | datastructures.ConstantData | datastructures.DagData | datastructures.ForEachData | datastructures.IfData @@ -70,16 +98,20 @@ def load( | datastructures.WhileData ): """Load a node or IO port using its lexical path.""" - return load_from_bag(self.bag, path) + return load_from_bag(self.bag, path, self.storage_root) @_import_alarm -def validate_bag(bag: H5Bag): +def validate_bag(bag: H5Bag, storage_root: str = DEFAULT_STORAGE_ROOT): if not isinstance(bag, boh.H5Bag): raise TypeError(f"Expected a {boh.H5Bag.__name__!r} object, got {bag!r}") _validate_bag_metadata(bag) - _validate_object_metadata(bag) + _validate_object_metadata(bag, storage_root) + + +def _normalize_root(storage_root: str) -> str: + return storage_root.rstrip("/") def _validate_bag_metadata(bag: H5Bag): @@ -103,23 +135,29 @@ def _validate_bag_metadata(bag: H5Bag): ) -def _validate_object_metadata(bag: H5Bag): - object_info = bag["object"] - if object_info.qualname != datastructures.DagData.__qualname__: +def _validate_object_metadata(bag: H5Bag, storage_root: str = DEFAULT_STORAGE_ROOT): + root = _normalize_root(storage_root) + try: + object_info = bag[root] + except boh.exceptions.InvalidMetadataError: + raise ValueError(f"Nothing is stored at {root!r}") from None + qualnames = tuple(cls.__qualname__ for cls in _NODE_DATA_TYPES) + if object_info.qualname not in qualnames: raise TypeError( - "Can only load saved workflow data " - f"({datastructures.DagData.__qualname__!r} type), but got " - f"{object_info.qualname!r}" + f"Can only load saved node data (one of {qualnames}), but got " + f"{object_info.qualname!r} at {root!r}" ) -def list_lexical_paths(bag: boh.H5Bag) -> list[str]: +def list_lexical_paths( + bag: boh.H5Bag, storage_root: str = DEFAULT_STORAGE_ROOT +) -> list[str]: """ Look through the bag and return a list of "."-separated lexical paths for nodes and - ports. + ports, starting from the node data stored at *storage_root*. """ paths: list[str] = [] - _collect_lexical_paths(bag, "object/", "", paths) + _collect_lexical_paths(bag, _normalize_root(storage_root), "", paths) return paths @@ -165,9 +203,10 @@ def _path_to_nodes(path: str) -> str: def load_from_bag( - bag: H5Bag, lexical_path: str + bag: H5Bag, lexical_path: str, storage_root: str = DEFAULT_STORAGE_ROOT ) -> ( datastructures.AtomicData + | datastructures.ConstantData | datastructures.DagData | datastructures.ForEachData | datastructures.IfData @@ -184,11 +223,13 @@ def load_from_bag( bag (H5Bag): The bag containing the saved node data. lexical_path (str): The dot-separated path of node names, IO references, and/or port names. + storage_root (str): Where in the bag the node data is kept. Defaults to the + top-level object. Returns: A retrospective data node or IO data port """ - storage_path = "object/" + storage_path = _normalize_root(storage_root) step = "" walked_path = step while lexical_path: @@ -210,19 +251,9 @@ def load_from_bag( f"from among {tuple(obj.keys())}" ) - expected_types = ( - datastructures.AtomicData, - datastructures.DagData, - datastructures.ForEachData, - datastructures.IfData, - datastructures.InputDataPort, - datastructures.OutputDataPort, - datastructures.TryData, - datastructures.WhileData, - ) - if not isinstance(obj, expected_types): + if not isinstance(obj, _LOADABLE_TYPES): raise TypeError( - f"Expected to load one of {tuple(cls.__name__ for cls in expected_types)}, " + f"Expected to load one of {tuple(cls.__name__ for cls in _LOADABLE_TYPES)}, " f"but got {type(obj).__name__}: {obj!r}" ) return obj diff --git a/src/flowrep/retrospective/storage_widget.py b/src/flowrep/retrospective/storage_widget.py index 14aa65b3..e5cc3fbc 100644 --- a/src/flowrep/retrospective/storage_widget.py +++ b/src/flowrep/retrospective/storage_widget.py @@ -63,7 +63,7 @@ def __init__(self, browser: storage.LexicalBagBrowser) -> None: root = self._make_node( label="Workflow", - storage_path="object/", + storage_path=browser.storage_root, lexical_path="", icon="project-diagram", opened=True, @@ -121,17 +121,24 @@ def _make_node( def _has_expandable_children(self, storage_path: str) -> bool: """True if storage_path has child nodes or IO port groups.""" - for suffix in ( - "/state/nodes", - "/state/input_ports", - "/state/output_ports", - ): - try: - if self._bag.open_group(f"{storage_path}{suffix}"): - return True - except KeyError: - continue - return False + return any( + self._open_group_or_empty(f"{storage_path}{suffix}") + for suffix in ( + "/state/nodes", + "/state/input_ports", + "/state/output_ports", + ) + ) + + def _open_group_or_empty(self, storage_path: str) -> list[str]: + """The members of the group at storage_path, or none if there is no group. + + Only composite node data has a group of child nodes. + """ + try: + return self._bag.open_group(storage_path) + except KeyError: + return [] def _lazy_expand(self, change: traitlets.Bunch) -> None: node = change["owner"] @@ -175,7 +182,7 @@ def _add_node_children(self, node: ipytree.Node) -> None: # Child nodes ------------------------------------------------------ nodes_storage = f"{storage_path}/state/nodes" - children = self._bag.open_group(nodes_storage) + children = self._open_group_or_empty(nodes_storage) for child in children: child_node = self._make_node( label=child, diff --git a/tests/flowrep_static/node_data.py b/tests/flowrep_static/node_data.py new file mode 100644 index 00000000..b88f97f2 --- /dev/null +++ b/tests/flowrep_static/node_data.py @@ -0,0 +1,67 @@ +"""Run instances of every concrete retrospective node data type, for storage tests.""" + +from flowrep import wfms, workflow +from flowrep.retrospective import datastructures + +from flowrep_static import library + + +@workflow +def looped(xs): + ys = [] + for x in xs: + y = library.increment(x) + ys.append(y) + return ys + + +@workflow +def branched(n): + # A statement, not a ternary: only an if-block parses into IfData + if library.is_positive(n): # noqa: SIM108 + m = library.increment(n) + else: + m = library.decrement(n) + return m + + +@workflow +def counted(n): + while library.is_positive(n): + n = library.decrement(n) + return n + + +@workflow +def guarded(x, y): + try: + z = library.raises_custom(x, y) + except library.MyCustomException: + z = library.increment(x) + return z + + +@workflow +def with_constant(a): + c = library.increment(a, 5) + return c + + +_RUNS = ( + (looped, {"xs": [1, 2]}), + (branched, {"n": 1}), + (counted, {"n": 2}), + (guarded, {"x": 1, "y": 2}), + (with_constant, {"a": 1}), +) + + +def samples() -> dict[type[datastructures.NodeData], datastructures.NodeData]: + """One freshly run instance of each node data type found in these workflows.""" + found: dict[type[datastructures.NodeData], datastructures.NodeData] = {} + for wf, kwargs in _RUNS: + data = wfms.run_recipe(wf.flowrep_recipe, **kwargs) + found.setdefault(type(data), data) + for child in data.nodes.values(): + found.setdefault(type(child), child) + return found diff --git a/tests/unit/test_storage.py b/tests/unit/test_storage.py index c2e510f4..f10c3ce6 100644 --- a/tests/unit/test_storage.py +++ b/tests/unit/test_storage.py @@ -3,6 +3,7 @@ and the LexicalBagBrowser convenience class). """ +import dataclasses import os import pathlib import tempfile @@ -12,7 +13,7 @@ from flowrep import wfms from flowrep.retrospective import datastructures, storage, storage_widget -from flowrep_static import library +from flowrep_static import library, node_data try: import bagofholding as boh @@ -356,5 +357,92 @@ def test_widget_returns_tree(self): self.assertIsInstance(result, storage_widget.LexicalBagTree) +# ═══════════════════════════════════════════════════════════════════════════ +# storage_root +# ═══════════════════════════════════════════════════════════════════════════ + + +@dataclasses.dataclass +class _Holder: + """Something that keeps node data inside itself, as a pyiron_workflow Run does.""" + + result: datastructures.NodeData + + +_HELD_ROOT = "object/state/result" + + +class TestStorageRoot(_BagTestCase): + def setUp(self) -> None: + super().setUp() + self.top_path = self._bag_path("top.h5") + self.held_path = self._bag_path("held.h5") + wf_data = _save_workflow(self.top_path, a=3, b=4) + boh.H5Bag.save(_Holder(wf_data), self.held_path) + + def test_held_data_lists_like_top_level_data(self): + self.assertEqual( + storage.list_lexical_paths(boh.H5Bag(self.top_path)), + storage.list_lexical_paths(boh.H5Bag(self.held_path), _HELD_ROOT), + ) + + def test_held_data_loads_like_top_level_data(self): + port = "typed_add_0.outputs.output_0" + self.assertEqual( + storage.load_from_bag(boh.H5Bag(self.top_path), port).value, + storage.load_from_bag(boh.H5Bag(self.held_path), port, _HELD_ROOT).value, + ) + + def test_browser_works_from_its_root(self): + browser = storage.LexicalBagBrowser(self.held_path, storage_root=_HELD_ROOT) + self.assertEqual(_HELD_ROOT, browser.storage_root) + self.assertIn("typed_add_0.outputs.output_0", browser.list_paths()) + self.assertIsInstance(browser.load("typed_add_0"), datastructures.AtomicData) + + def test_a_trailing_slash_is_tolerated(self): + browser = storage.LexicalBagBrowser( + self.held_path, storage_root=f"{_HELD_ROOT}/" + ) + self.assertEqual(_HELD_ROOT, browser.storage_root) + self.assertIn("inputs.a", browser.list_paths()) + + def test_the_default_root_rejects_a_holder(self): + with self.assertRaisesRegex(TypeError, _Holder.__qualname__): + storage.LexicalBagBrowser(self.held_path) + + def test_an_empty_root_raises(self): + with self.assertRaisesRegex(ValueError, "object/state/nothing"): + storage.LexicalBagBrowser( + self.held_path, storage_root="object/state/nothing" + ) + + +class TestEveryNodeDataType(_BagTestCase): + """Each kind of node data, saved as the bag's top-level object.""" + + def setUp(self) -> None: + super().setUp() + self.samples = node_data.samples() + + def test_samples_cover_every_node_data_type(self): + self.assertSetEqual(set(storage._NODE_DATA_TYPES), set(self.samples)) + + def test_each_type_validates_lists_and_loads(self): + for cls, data in self.samples.items(): + with self.subTest(cls=cls.__name__): + path = self._bag_path(f"{cls.__name__}.h5") + boh.H5Bag.save(data, path) + browser = storage.LexicalBagBrowser(path) + paths = browser.list_paths() + self.assertTrue( + {f"outputs.{port}" for port in data.output_ports}.issubset(paths) + ) + self.assertIsInstance(browser.load(""), cls) + for lexical_path in paths: + self.assertIsInstance( + browser.load(lexical_path), storage._LOADABLE_TYPES + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_storage_widget.py b/tests/unit/test_storage_widget.py index 542cad99..e3180e73 100644 --- a/tests/unit/test_storage_widget.py +++ b/tests/unit/test_storage_widget.py @@ -15,7 +15,7 @@ storage_widget as sw, ) -from flowrep_static import library +from flowrep_static import library, node_data try: import bagofholding as boh @@ -113,7 +113,18 @@ def test_root_metadata_registered(self): with self.subTest(tree=tree): meta = tree._meta(self._get_root_node(tree)) self.assertEqual(meta.lexical_path, "") - self.assertEqual(meta.storage_path, "object/") + self.assertEqual(meta.storage_path, self.browser.storage_root) + + def test_root_follows_the_browser_storage_root(self): + self.browser.storage_root = "object/elsewhere" + with mock.patch.object( + storage_widget.LexicalBagTree, + "_has_expandable_children", + return_value=False, + ): + tree = storage_widget.LexicalBagTree(self.browser) + root = self._get_root_node(tree) + self.assertEqual("object/elsewhere", tree._meta(root).storage_path) def test_root_is_opened_and_loaded(self): for tree in (self.tree, self.void_tree): @@ -319,5 +330,36 @@ def test_port_lexical_path(self): self.assertEqual(meta.lexical_path, "inputs.a") +@unittest.skipUnless(_has_ipytree, "ipytree not installed") +class TestEveryNodeDataType(unittest.TestCase): + """Each kind of node data, saved as the bag's top-level object.""" + + def setUp(self) -> None: + self._tmpdir = tempfile.TemporaryDirectory() + self.addCleanup(self._tmpdir.cleanup) + + @staticmethod + def _expand_all(tree: "storage_widget.LexicalBagTree", node) -> None: + node.opened = True + for child in node.nodes: + if not child.disabled: + TestEveryNodeDataType._expand_all(tree, child) + + def test_fully_expanded_tree_shows_what_storage_lists(self): + for cls, data in node_data.samples().items(): + with self.subTest(cls=cls.__name__): + path = os.path.join(self._tmpdir.name, f"{cls.__name__}.h5") + boh.H5Bag.save(data, path) + browser = storage.LexicalBagBrowser(path) + tree = storage_widget.LexicalBagTree(browser) + self._expand_all(tree, tree.nodes[0]) + shown = { + meta.lexical_path + for meta in tree._node_meta.values() + if meta.lexical_path and not meta.is_io_group + } + self.assertSetEqual(set(browser.list_paths()), shown) + + if __name__ == "__main__": unittest.main()