Supporting pruning and merging with the new branch types (#5)

Updated pruning and merging logic to preserve correctness in control flow graphs containing except_branches and reraise_branches.

python_graphs/control_flow.py

@@ -365,6 +365,12 @@ def remove_exit(self, block):
     for branch_decision, branch_exit in self.branches.copy().items():
       if branch_exit is block:
         del self.branches[branch_decision]
+    for branch_decision, branch_exit in self.except_branches.copy().items():
+      if branch_exit is block:
+        del self.except_branches[branch_decision]
+    for branch_decision, branch_exit in self.reraise_branches.copy().items():
+      if branch_exit is block:
+        del self.reraise_branches[branch_decision]
 
   def can_prune(self):
     return self.is_empty() and self.prunable
@@ -385,6 +391,8 @@ def prune(self):
       exits_from_middle = prev_block.exits_from_middle.copy()
       exits_from_end = prev_block.exits_from_end.copy()
       branches = prev_block.branches.copy()
+      except_branches = prev_block.except_branches.copy()
+      reraise_branches = prev_block.reraise_branches.copy()
       for next_block in nexts:
         if self in exits_from_middle:
           prev_block.add_exit(next_block, interrupting=True)
@@ -394,6 +402,12 @@ def prune(self):
         for branch_decision, branch_exit in branches.items():
           if branch_exit is self:
             prev_block.branches[branch_decision] = next_block
+        for branch_decision, branch_exit in except_branches.items():
+          if branch_exit is self:
+            prev_block.except_branches[branch_decision] = next_block
+        for branch_decision, branch_exit in reraise_branches.items():
+          if branch_exit is self:
+            prev_block.reraise_branches[branch_decision] = next_block
 
     for prev_block in prevs:
       prev_block.remove_exit(self)
@@ -427,6 +441,10 @@ def merge(self):
 
     for branch_decision, branch_exit in next_block.branches.items():
       self.branches[branch_decision] = branch_exit
+    for branch_decision, branch_exit in next_block.except_branches.items():
+      self.except_branches[branch_decision] = branch_exit
+    for branch_decision, branch_exit in next_block.reraise_branches.items():
+      self.reraise_branches[branch_decision] = branch_exit
 
     self.remove_exit(next_block)
     for block in next_block.next.copy():

setup.py

@@ -38,7 +38,7 @@
 TEST_DEPENDENCIES = [
 ]
 
-VERSION = '1.2.0'
+VERSION = '1.2.1'
 URL = 'https://github.com/google-research/python-graphs'
 
 setup(