Export to dict

pyiron_workflow/nodes/composite.py

@@ -13,6 +13,7 @@
 from pyiron_snippets.colors import SeabornColors
 from pyiron_snippets.dotdict import DotDict
 
+from pyiron_workflow.channels import NOT_DATA, Channel
 from pyiron_workflow.create import HasCreator
 from pyiron_workflow.mixin.semantics import SemanticParent
 from pyiron_workflow.node import Node
@@ -26,6 +27,130 @@
     from pyiron_workflow.storage import StorageInterface
 
 
+def _extract_data(item: Channel, with_values=True, with_default=True) -> dict:
+    data = {}
+    data_dict = {"default": NOT_DATA, "value": NOT_DATA, "type_hint": None}
+    if not with_values:
+        data_dict.pop("value")
+    if not with_default:
+        data_dict.pop("default")
+    for key, value in data_dict.items():
+        if getattr(item, key) is not value:
+            data[key] = getattr(item, key)
+    return data
+
+
+def _is_internal_connection(channel: Channel, workflow: Composite, io_: str) -> bool:
+    """
+    Check if a channel is connected to another channel in the same workflow.
+
+    Args:
+        channel (Channel): The channel to check.
+        workflow (Composite): The workflow to check whether the channel is connected to.
+        io_ (str): The IO direction to check.
+
+    Returns:
+        bool: Whether the channel is connected to another channel in the same workflow.
+    """
+    if not channel.connected:
+        return False
+    return any(channel.connections[0] in getattr(n, io_) for n in workflow)
+
+
+def _get_scoped_label(channel: Channel, io_: str) -> str:
+    return channel.scoped_label.replace("__", f".{io_}.")
+
+
+def _io_to_dict(
+    node: Node, with_values: bool = True, with_default: bool = True
+) -> dict:
+    data = {"inputs": {}, "outputs": {}}
+    is_composite = isinstance(node, Composite)
+    for io_ in ["inputs", "outputs"]:
+        for inp in getattr(node, io_):
+            if is_composite:
+                data[io_][inp.scoped_label] = _extract_data(
+                    inp, with_values=with_values, with_default=with_default
+                )
+            else:
+                data[io_][inp.label] = _extract_data(
+                    inp, with_values=with_values, with_default=with_default
+                )
+    return data
+
+
+def export_node_to_dict(
+    node: Node, with_values: bool = True, with_default: bool = True
+) -> dict:
+    """
+    Export a node to a dictionary.
+
+    Args:
+        node (Node): The node to export.
+        with_values (bool): Whether to include the values of the channels in the
+            dictionary. (Default is True.)
+
+    Returns:
+        dict: The exported node as a dictionary.
+    """
+    data = {"inputs": {}, "outputs": {}, "function": node.node_function}
+    data.update(_io_to_dict(node, with_values=with_values, with_default=with_default))
+    return data
+
+
+def export_composite_to_dict(
+    workflow: Composite, with_values: bool = True, with_default: bool = True
+) -> dict:
+    """
+    Export a composite to a dictionary.
+
+    Args:
+        workflow (Composite): The composite to export.
+        with_values (bool): Whether to include the values of the channels in the
+            dictionary. (Default is True.)
+
+    Returns:
+        dict: The exported composite as a dictionary.
+    """
+    data = {"inputs": {}, "outputs": {}, "nodes": {}, "edges": []}
+    for inp in workflow.inputs:
+        if inp.value_receiver is not None:
+            data["edges"].append(
+                (
+                    f"inputs.{inp.scoped_label}",
+                    _get_scoped_label(inp.value_receiver, "inputs"),
+                )
+            )
+    for node in workflow:
+        label = node.label
+        if isinstance(node, Composite):
+            data["nodes"][label] = export_composite_to_dict(
+                node, with_values=with_values
+            )
+        else:
+            data["nodes"][label] = export_node_to_dict(node, with_values=with_values)
+        for inp in node.inputs:
+            if _is_internal_connection(inp, workflow, "outputs"):
+                data["edges"].append(
+                    (
+                        _get_scoped_label(inp.connections[0], "outputs"),
+                        _get_scoped_label(inp, "inputs"),
+                    )
+                )
+        for out in node.outputs:
+            if out.value_receiver is not None:
+                data["edges"].append(
+                    (
+                        _get_scoped_label(out, "outputs"),
+                        f"outputs.{out.value_receiver.scoped_label}",
+                    )
+                )
+    data.update(
+        _io_to_dict(workflow, with_values=with_values, with_default=with_default)
+    )
+    return data
+
+
 def _get_graph_as_dict(composite: Composite) -> dict:
     if not isinstance(composite, Composite):
         return composite
@@ -459,6 +584,13 @@ def graph_as_dict(self) -> dict:
         """
         return _get_graph_as_dict(self)
 
+    def export_to_dict(
+        self, with_values: bool = True, with_default: bool = True
+    ) -> dict:
+        return export_composite_to_dict(
+            self, with_values=with_values, with_default=with_default
+        )
+
     def _get_connections_as_strings(
         self, panel_getter: Callable
     ) -> list[tuple[tuple[str, str], tuple[str, str]]]:

tests/unit/nodes/test_export.py

@@ -0,0 +1,44 @@
+import unittest
+
+from pyiron_workflow._tests import ensure_tests_in_python_path
+from pyiron_workflow.nodes.composite import Composite
+from pyiron_workflow.workflow import Workflow
+
+ensure_tests_in_python_path()
+
+
+@Workflow.wrap.as_function_node
+def add_one(a: int):
+    result = a + 1
+    return result
+
+@Workflow.wrap.as_function_node
+def add_two(b: int = 10) -> int:
+    result = b + 2
+    return result
+
+@Workflow.wrap.as_macro_node
+def add_three(macro: Composite | None = None, c: int = 0) -> int:
+    macro.one = add_one(a=c)
+    macro.two = add_two(b=macro.one)
+    w = macro.two
+    return w
+
+
+class TestExport(unittest.TestCase):
+    def test_io_independence(self):
+        wf = Workflow("my_wf")
+        wf.three = add_three(c=1)
+        wf.four = add_one(a=wf.three)
+        wf.run()
+        data = wf.export_to_dict()
+        self.assertEqual(
+            set(data.keys()), {"edges", "inputs", "nodes", "outputs"}
+        )
+        self.assertEqual(
+            data["inputs"], {'three__c': {"default": 0, 'value': 1, 'type_hint': int}}
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()