Commit edfddeac authored by Wout De Nolf's avatar Wout De Nolf
Browse files

esrftaskgraph: refactor extract_subgraphs

parent 923564a1
......@@ -125,6 +125,20 @@ def get_subgraphs(graph):
return subgraphs
def _pop_subgraph_node_name(subgraph_name, link_attrs, source=True):
if source:
key = "source"
else:
key = "target"
try:
subgraph_node_name = link_attrs.pop(key)
except KeyError:
raise ValueError(
f"The '{key}' attribute to specify a node in subgraph '{subgraph_name}' is missing"
) from None
return subgraph_name, subgraph_node_name
def extract_subgraphs(graph, subgraphs):
# Edges between supergraph and subgraphs
edges = list()
......@@ -135,23 +149,37 @@ def extract_subgraphs(graph, subgraphs):
links = graph[source_name][target_name].get("links", list())
for link_attrs in links:
link_attrs = dict(link_attrs)
source = link_attrs.pop("source", source_name)
if source_name in subgraphs:
source = (source_name, source)
target = (target_name, link_attrs.pop("target"))
source = _pop_subgraph_node_name(
source_name, link_attrs, source=True
)
else:
link_attrs.pop("source", None)
source = source_name
target = _pop_subgraph_node_name(target_name, link_attrs, source=False)
update_attrs[target] = link_attrs.pop("node_attributes", dict())
edges.append((source, target, link_attrs))
for target_name in graph.successors(subgraph_name):
source_name = subgraph_name
links = graph[source_name][target_name].get("links", list())
for link_attrs in links:
link_attrs = dict(link_attrs)
target = link_attrs.pop("target", target_name)
if target_name in subgraphs:
target = (target_name, target)
source = (source_name, link_attrs.pop("source"))
target = _pop_subgraph_node_name(
target_name, link_attrs, source=False
)
else:
link_attrs.pop("target", None)
target = target_name
source = _pop_subgraph_node_name(source_name, link_attrs, source=True)
link_attrs.pop("node_attributes", None)
edges.append((source, target, link_attrs))
graph.remove_nodes_from(subgraphs.keys())
return edges, update_attrs
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment