esrftaskgraph: refactor extract_subgraphs

esrftaskgraph/esrftaskgraph/graph.py

@@ -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