Refactored eades append/prepend

src/macrogen/fes.py

@@ -14,6 +14,17 @@
 
 class Eades:
 
+    def to_start(self, node):
+        """
+        Removes the node from the graph and appends it to the start sequence.
+        """
+        self.start.append(node)
+        self.graph.remove_node(node)
+
+    def to_end(self, node):
+        self.end.insert(0, node)
+        self.graph.remove_node(node)
+
     def _exhaust_sinks(self, sink: bool = True):
         """
         Produces all sinks until there are no more.
@@ -25,7 +36,6 @@ def _exhaust_sinks(self, sink: bool = True):
             sinks = [u for (u, d) in sink_method() if d == 0]
             if sinks:
                 yield from sinks
-                self.graph.remove_nodes_from(sinks)
             else:
                 return
 
@@ -77,14 +87,13 @@ def solve(self) -> List[Tuple[V, V]]:
         self.graph.remove_edges_from(list(self.graph.selfloop_edges()))
         while self.graph:
             for v in self._exhaust_sinks():
-                self.end.insert(0, v)
+                self.to_end(v)
             for v in self._exhaust_sources():
-                self.start.append(v)
+                self.to_start(v)
             if self.graph:
                 u = max(self.graph.nodes, key=lambda v: self.graph.out_degree(v, weight='weight')
                                                         - self.graph.in_degree(v, weight='weight'))
-                self.start.append(u)
-                self.graph.remove_node(u)
+                self.to_start(u)
         ordering = self.start + self.end
         pos = dict(zip(ordering, itertools.count()))
         feedback_edges = list(self.original_graph.selfloop_edges())

tests/test_fes.py

@@ -19,14 +19,12 @@ def graph1():
 
 def test_all_sinks(graph1):
     eades = Eades(graph1)
-    sinks = list(eades._exhaust_sinks())
+    sinks = []
+    for sink in eades._exhaust_sinks():
+        eades.graph.remove_node(sink)
+        sinks.append(sink)
     assert sinks == [5, 4]
 
-
-def test_all_sources(graph1):
-    assert list(Eades(graph1)._exhaust_sources()) == [1]
-
-
 def test_eades(graph1):
     assert list(eades(graph1)) == [(3, 2)]