Tighten flow charts more

ceci/pipeline/graph.py

@@ -207,13 +207,13 @@ def combine_output_nodes(graph):
 
     Parameters
     ----------
-    graph : networkx.DiGraph
+    graph : pygraphviz.AGraph
         The pipeline graph to modify
     """
     groups = collections.defaultdict(list)
     # find groups of nodes that have no outgoing edges and all
     # share the same incoming edge.
-    for node in graph.nodes_iter():
+    for node in graph.nodes():
         node_data = graph.get_node(node).attr
         if node_data["type"] != "output" or len(graph.successors(node)) != 0:
             continue
@@ -238,3 +238,40 @@ def combine_output_nodes(graph):
         graph.add_edge(parent, new_node_name)
         # remove the old nodes
         graph.remove_nodes_from(group)
+
+
+def combine_input_nodes(graph):
+    """
+    Combine input nodes that share the same successor and have no
+    predecessors into a single node.
+
+    This is useful for plotting flow charts.
+
+    Parameters
+    ----------
+    graph : pygraphviz.AGraph
+        The pipeline graph to modify
+    """
+    groups = collections.defaultdict(list)
+    for node in graph.nodes():
+        node_data = graph.get_node(node).attr
+        if node_data["type"] != "input" or len(graph.predecessors(node)) != 0:
+            continue
+        successors = list(graph.successors(node))
+
+        # This shouldn't happen in the default case,
+        # but might if things change later, so let's deal with it
+        # just in case.
+        if len(successors) != 1:
+            continue
+
+        successor = successors[0]
+        groups[successor].append(node)
+
+    for successor, group in groups.items():
+        if len(group) < 2:
+            continue
+        new_node_name = "\n".join(group)
+        graph.add_node(new_node_name, type="input")
+        graph.add_edge(new_node_name, successor)
+        graph.remove_nodes_from(group)

ceci/pipeline/pipeline.py

@@ -9,7 +9,7 @@
 from ..stage import PipelineStage
 from ..sites import load, set_default_site, get_default_site
 from ..utils import extra_paths
-from .graph import build_graph, get_static_ordering, trim_pipeline_graph, combine_output_nodes
+from .graph import build_graph, get_static_ordering, trim_pipeline_graph, combine_output_nodes, combine_input_nodes
 from .file_manager import FileManager
 from .sec import StageExecutionConfig
 from .templates import read_and_apply_template
@@ -910,9 +910,10 @@ def make_flow_chart(self, filename=None):
         # are final outputs not used in the rest of the pipeline. We
         # can combine these to make the graph more readable.
         combine_output_nodes(graph)
+        combine_input_nodes(graph)
 
         # set the colours and styles for the boxes
-        for node in graph.nodes_iter():
+        for node in graph.nodes():
             node_type = node.attr['type']
             if node_type == "input":
                 node.attr.update(shape="box", color="gold", style="filled")

tests/test_graph.py

@@ -1,5 +1,6 @@
-from ceci.pipeline.graph import trim_pipeline_graph
+from ceci.pipeline.graph import trim_pipeline_graph, combine_input_nodes
 import networkx
+import pygraphviz
 
 
 def check_membership(graph, everything, expected):
@@ -97,3 +98,23 @@ def test_trim_pipeline_graph():
     }
 
     check_membership(sub4, everything, expected4)
+
+
+def test_combine_input_nodes():
+    graph = pygraphviz.AGraph(directed=True)
+    graph.add_node("a", type="input")
+    graph.add_node("b", type="input")
+    graph.add_node("stage", type="stage")
+    graph.add_edge("a", "stage")
+    graph.add_edge("b", "stage")
+
+    combine_input_nodes(graph)
+
+    input_nodes = [
+        str(node)
+        for node in graph.nodes()
+        if graph.get_node(node).attr["type"] == "input"
+    ]
+    assert len(input_nodes) == 1
+    assert input_nodes[0] == "a\nb"
+    assert [str(node) for node in graph.predecessors("stage")] == ["a\nb"]