Commits

Vinay Sajip committed 8a1f82c

Enhanced sequencer removal functionality.

  • Participants
  • Parent commits 4621e02

Comments (0)

Files changed (2)

File distlib/util.py

     def add_node(self, node):
         self._nodes.add(node)
 
-    def remove_node(self, node):
-        self._nodes.remove(node)
+    def remove_node(self, node, edges=False):
+        if node in self._nodes:
+            self._nodes.remove(node)
+        if edges:
+            for p in set(self._preds.get(node, ())):
+                self.remove(p, node)
+            for s in set(self._succs.get(node, ())):
+                self.remove(node, s)
+            # Remove empties
+            for k, v in list(self._preds.items()):
+                if not v:
+                    del self._preds[k]
+            for k, v in list(self._succs.items()):
+                if not v:
+                    del self._succs[k]
 
     def add(self, pred, succ):
         assert pred != succ

File tests/test_util.py

         seq.remove('C', 'A')
         self.assertEqual(list(seq.get_steps('D')), ['A', 'B', 'C', 'D'])
 
+    def test_sequencer_removal(self):
+        seq = Sequencer()
+        seq.add('A', 'B')
+        seq.add('B', 'C')
+        seq.add('C', 'D')
+        preds = {
+            'B': set(['A']),
+            'C': set(['B']),
+            'D': set(['C'])
+        }
+        succs =  {
+            'A': set(['B']),
+            'B': set(['C']),
+            'C': set(['D'])
+        }
+        self.assertEqual(seq._preds, preds)
+        self.assertEqual(seq._succs, succs)
+        seq.remove_node('C')
+        self.assertEqual(seq._preds, preds)
+        self.assertEqual(seq._succs, succs)
+        seq.remove_node('C', True)
+        self.assertEqual(seq._preds, {'B': set(['A'])})
+        self.assertEqual(seq._succs, {'A': set(['B'])})
+
     def test_unarchive(self):
         import zipfile, tarfile