diff --git a/src/ewoksdraw/graph_to_svg.py b/src/ewoksdraw/graph_to_svg.py index e91a3b6..bd49bc5 100644 --- a/src/ewoksdraw/graph_to_svg.py +++ b/src/ewoksdraw/graph_to_svg.py @@ -13,17 +13,32 @@ from .svg.svg_canvas import SvgCanvas from .svg.svg_task import SvgTask from .svg.svg_task_group import SvgTaskGroup +from .utils import get_edge_sources_and_targets def build_svg_task_group(graph: TaskGraph) -> SvgTaskGroup: """Build an SVG task group from an Ewoks task graph.""" + source_outputs_per_node, target_inputs_per_node = get_edge_sources_and_targets( + graph + ) + svg_tasks = {} for node_id, node_attrs in graph.graph.nodes.items(): signature = node_signature(node_id, node_attrs) + + # Use dict instead of set to remove duplicate while keeping order + inputs = dict.fromkeys(node_input.name for node_input in signature.inputs) + # Add eventual missing target inputs + inputs.update(dict.fromkeys(target_inputs_per_node[node_id])) + + outputs = dict.fromkeys(node_output.name for node_output in signature.outputs) + # Add eventual missing source outputs + outputs.update(dict.fromkeys(source_outputs_per_node[node_id])) + svg_tasks[node_id] = SvgTask( task_name=node_id, - input_names=[node_input.name for node_input in signature.inputs], - output_names=[node_output.name for node_output in signature.outputs], + input_names=list(inputs.keys()), + output_names=list(outputs.keys()), import_error=bool(signature.import_error), ) return SvgTaskGroup( diff --git a/src/ewoksdraw/layout/elk_converter.py b/src/ewoksdraw/layout/elk_converter.py index 78cf9db..4c4b2ca 100644 --- a/src/ewoksdraw/layout/elk_converter.py +++ b/src/ewoksdraw/layout/elk_converter.py @@ -1,6 +1,7 @@ -import warnings from typing import Any +from ..utils import get_valid_data_mappings + try: from typing import NotRequired except ImportError: @@ -152,38 +153,19 @@ def convert_ewoks_to_elk_graph( used_ids.add(root_id) edges: list[ElkEdgeBeforeLayout] = [] - for source, target, link_attrs in ewoks_graph.graph.edges(data=True): - if link_attrs.get("map_all_data", False): - warnings.warn( - f"Ewoks link {source!r} -> {target!r} uses 'map_all_data', which " - "is not yet supported.", - UserWarning, - stacklevel=2, - ) - - for mapping in link_attrs.get("data_mapping", []): - source_output = mapping.get("source_output") - if source_output is None: - warnings.warn( - f"Data mapping on Ewoks link {source!r} -> {target!r} has no " - "'source_output', which is not yet supported.", - UserWarning, - stacklevel=2, - ) - continue - - edge_id = _available_elk_id( - f"edge_{len(edges)}_{source}_{target}", used_ids - ) - - edges.append( - { - "id": edge_id, - "sources": [f"{source}.output.{source_output}"], - "targets": [f"{target}.input.{mapping['target_input']}"], - } - ) - used_ids.add(edge_id) + for mapping in get_valid_data_mappings(ewoks_graph): + edge_id = _available_elk_id( + f"edge_{len(edges)}_{mapping.source}_{mapping.target}", used_ids + ) + + edges.append( + { + "id": edge_id, + "sources": [f"{mapping.source}.output.{mapping.source_output}"], + "targets": [f"{mapping.target}.input.{mapping.target_input}"], + } + ) + used_ids.add(edge_id) return { "id": root_id, diff --git a/src/ewoksdraw/models.py b/src/ewoksdraw/models.py new file mode 100644 index 0000000..9c2cd3d --- /dev/null +++ b/src/ewoksdraw/models.py @@ -0,0 +1,8 @@ +from typing import NamedTuple + + +class DataMapping(NamedTuple): + source: str + target: str + source_output: str + target_input: str diff --git a/src/ewoksdraw/utils.py b/src/ewoksdraw/utils.py new file mode 100644 index 0000000..8dbe678 --- /dev/null +++ b/src/ewoksdraw/utils.py @@ -0,0 +1,53 @@ +import warnings +from collections import defaultdict + +from ewokscore.graph import TaskGraph + +from .models import DataMapping + + +def get_valid_data_mappings(ewoks_graph: TaskGraph) -> list[DataMapping]: + data_mappings: list[DataMapping] = [] + for source, target, link_attrs in ewoks_graph.graph.edges(data=True): + if link_attrs.get("map_all_data", False): + warnings.warn( + f"Ewoks link {source!r} -> {target!r} uses 'map_all_data', which " + "is not yet supported.", + UserWarning, + stacklevel=2, + ) + + for mapping in link_attrs.get("data_mapping", []): + source_output = mapping.get("source_output") + if source_output is None: + warnings.warn( + f"Data mapping on Ewoks link {source!r} -> {target!r} has no " + "'source_output', which is not yet supported.", + UserWarning, + stacklevel=2, + ) + continue + + data_mappings.append( + DataMapping( + source=source, + target=target, + source_output=source_output, + target_input=mapping["target_input"], + ) + ) + + return data_mappings + + +def get_edge_sources_and_targets( + graph: TaskGraph, +) -> tuple[dict[str, list[str]], dict[str, list[str]]]: + data_mappings = get_valid_data_mappings(graph) + sources: dict[str, list[str]] = defaultdict(list[str]) + targets: dict[str, list[str]] = defaultdict(list[str]) + for mapping in data_mappings: + sources[mapping.source].append(mapping.source_output) + targets[mapping.target].append(mapping.target_input) + + return sources, targets