115 lines
3.5 KiB
Python
115 lines
3.5 KiB
Python
"""Shared workflow node traversal helpers."""
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, Iterator, Optional
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class WorkflowNodeContext:
|
|
"""A workflow node together with the scope where it is defined."""
|
|
|
|
node: Any
|
|
subgraph_id: Optional[Any] = ""
|
|
subgraph_name: Optional[Any] = ""
|
|
is_top_level: bool = True
|
|
|
|
|
|
def iter_workflow_nodes_with_scope(
|
|
workflow_json: Dict[str, Any],
|
|
) -> Iterator[WorkflowNodeContext]:
|
|
"""Yield top-level and subgraph nodes in serialized workflow order."""
|
|
nodes = workflow_json.get("nodes", [])
|
|
if isinstance(nodes, list):
|
|
for node in nodes:
|
|
yield WorkflowNodeContext(node=node)
|
|
|
|
definitions = workflow_json.get("definitions", {})
|
|
subgraphs = (
|
|
definitions.get("subgraphs", [])
|
|
if isinstance(definitions, dict)
|
|
else []
|
|
)
|
|
if not isinstance(subgraphs, list):
|
|
return
|
|
|
|
for subgraph in subgraphs:
|
|
if not isinstance(subgraph, dict):
|
|
continue
|
|
subgraph_nodes = subgraph.get("nodes", [])
|
|
if not isinstance(subgraph_nodes, list):
|
|
continue
|
|
|
|
subgraph_id = subgraph.get("id") or ""
|
|
subgraph_name = subgraph.get("name") or subgraph_id or ""
|
|
for node in subgraph_nodes:
|
|
yield WorkflowNodeContext(
|
|
node=node,
|
|
subgraph_id=subgraph_id,
|
|
subgraph_name=subgraph_name,
|
|
is_top_level=False,
|
|
)
|
|
|
|
|
|
def iter_active_workflow_nodes_with_scope(
|
|
workflow_json: Dict[str, Any],
|
|
) -> Iterator[WorkflowNodeContext]:
|
|
"""Yield top-level nodes and recursively referenced subgraph nodes."""
|
|
if not isinstance(workflow_json, dict):
|
|
return
|
|
|
|
top_level_nodes = workflow_json.get("nodes")
|
|
if not isinstance(top_level_nodes, list):
|
|
return
|
|
|
|
active_nodes = [node for node in top_level_nodes if isinstance(node, dict)]
|
|
for node in active_nodes:
|
|
yield WorkflowNodeContext(node=node)
|
|
|
|
definitions = workflow_json.get("definitions")
|
|
if not isinstance(definitions, dict):
|
|
return
|
|
|
|
subgraph_list = definitions.get("subgraphs")
|
|
if not isinstance(subgraph_list, list):
|
|
return
|
|
|
|
subgraphs = {
|
|
str(subgraph.get("id")): subgraph
|
|
for subgraph in subgraph_list
|
|
if isinstance(subgraph, dict) and subgraph.get("id") is not None
|
|
}
|
|
pending_subgraphs = [
|
|
str(node.get("type"))
|
|
for node in active_nodes
|
|
if str(node.get("type")) in subgraphs
|
|
]
|
|
seen_subgraphs = set()
|
|
|
|
while pending_subgraphs:
|
|
subgraph_id = pending_subgraphs.pop(0)
|
|
if subgraph_id in seen_subgraphs:
|
|
continue
|
|
seen_subgraphs.add(subgraph_id)
|
|
|
|
subgraph = subgraphs.get(subgraph_id)
|
|
subgraph_nodes = subgraph.get("nodes") if isinstance(subgraph, dict) else None
|
|
if not isinstance(subgraph_nodes, list):
|
|
continue
|
|
|
|
subgraph_name = subgraph.get("name") or subgraph_id or ""
|
|
for node in subgraph_nodes:
|
|
if not isinstance(node, dict):
|
|
continue
|
|
yield WorkflowNodeContext(
|
|
node=node,
|
|
subgraph_id=subgraph_id,
|
|
subgraph_name=subgraph_name,
|
|
is_top_level=False,
|
|
)
|
|
nested_subgraph_id = str(node.get("type"))
|
|
if (
|
|
nested_subgraph_id in subgraphs
|
|
and nested_subgraph_id not in seen_subgraphs
|
|
):
|
|
pending_subgraphs.append(nested_subgraph_id)
|