Add module frame to catch raises in top-level code, and introduces 3 …

…branch kinds. (#2)

1. Adds a module frame to catch raises in top-level code.
2. Also marks except expressions and finally blocks as branch points.

This includes PR #3: Splits "branches" into branches, except_branches, and reraise_branches.
* branches are your usual branch decisions: ifs, fors, and whiles.
* except_branches are at "except E:" statements, with True indicating the exception matches and False indicating it does not
* reraise_branches are at the end of "finally:" blocks, with True indicating the path taken after finally if an error has been raised previously, and False indicating the path taken if there's nothing to reraise at the end of the finally.

.gitignore

@@ -0,0 +1,9 @@
+scratch
+
+.DS_Store
+__MACOSX
+
+*~
+__pycache__
+.pytest_cache
+python_graphs.egg-info

python_graphs/control_flow.py

@@ -63,6 +63,10 @@ def new_block(self, node=None, label=None, prunable=True):
     self.blocks.append(block)
     return block
 
+  def move_block_to_rear(self, block):
+    self.blocks.remove(block)
+    self.blocks.append(block)
+
   def get_control_flow_nodes(self):
     return self.nodes
 
@@ -240,6 +244,7 @@ class Frame(object):
   """
 
   # Kinds:
+  MODULE = 'module'
   LOOP = 'loop'
   FUNCTION = 'function'
   TRY_EXCEPT = 'try-except'
@@ -295,6 +300,9 @@ def __init__(self, node=None, label=None, prunable=True):
     self.control_flow_node_indexes = None
 
     self.branches = {}
+    self.except_branches = {}
+    self.reraise_branches = {}
+
     self.exits_from_middle = set()
     self.exits_from_end = set()
     self.node = node
@@ -327,13 +335,18 @@ def raises_to(self, block):
     """Whether this block exits to `block` in the case of an exception."""
     return block in self.next and block in self.exits_from_middle
 
-  def add_exit(self, block, interrupting=False, branch=None):
+  def add_exit(self, block, interrupting=False,
+               branch=None, except_branch=None, reraise_branch=None):
     """Adds an exit from this block to `block`."""
     self.next.add(block)
     block.prev.add(self)
 
     if branch is not None:
       self.branches[branch] = block
+    if except_branch is not None:
+      self.except_branches[except_branch] = block
+    if reraise_branch is not None:
+      self.reraise_branches[reraise_branch] = block
 
     if interrupting:
       self.exits_from_middle.add(block)
@@ -612,8 +625,6 @@ def __init__(self):
   def run(self, node):
     start_block = self.graph.start_block
     end_block = self.visit(node, start_block)
-    exit_block = self.new_block(node=node, label='<exit>', prunable=False)
-    end_block.add_exit(exit_block)
     self.graph.compact()
 
   def visit(self, node, current_block):
@@ -658,7 +669,7 @@ def add_instruction(self, block, instruction):
     if not block.exits_from_middle:
       self.raise_through_frames(block, interrupting=True)
 
-  def raise_through_frames(self, block, interrupting=True):
+  def raise_through_frames(self, block, interrupting=True, except_branch=None):
     """Adds exits for the control flow of a raised exception.
 
     `interrupting` means the exit can occur at any point (exit_from_middle).
@@ -672,31 +683,52 @@ def raise_through_frames(self, block, interrupting=True):
       block: The block where the exception's control flow begins.
       interrupting: Whether the exception can be raised from any point in block.
         If False, the exception is only raised from the end of block.
+      except_branch: False indicates the node raising is doing so the because an exception
+        header did not match the raised error. None indicates otherwise.
     """
     frames = self.get_current_exception_handling_frames()
 
     if frames is None:
       return
 
+    # reraise_branch indicates whether the a raise is a reraise of an earlier exception.
+    # This is True after raising through a finally block, and None otherwise.
+    reraise_branch = None
+
     for frame in frames:
       if frame.kind == Frame.TRY_FINALLY:
         # Exit to finally and have finally exit to whatever's next...
         final_block = frame.blocks['final_block']
-        block.add_exit(final_block, interrupting=interrupting)
+        block.add_exit(final_block, interrupting=interrupting, except_branch=except_branch, reraise_branch=reraise_branch)
         block = frame.blocks['final_block_end']
         interrupting = False
+        # "True" indicates the path taken after finally if an error has been raised.
+        except_branch = None
+        reraise_branch = True
       elif frame.kind == Frame.TRY_EXCEPT:
         handler_block = frame.blocks['handler_block']
-        block.add_exit(handler_block, interrupting=interrupting)
-        interrupting = False  # return...
+        block.add_exit(handler_block, interrupting=interrupting, except_branch=except_branch, reraise_branch=reraise_branch)
+        # This will be the last frame in frames.
       elif frame.kind == Frame.FUNCTION:
         raise_block = frame.blocks['raise_block']
-        block.add_exit(raise_block, interrupting=interrupting)
+        block.add_exit(raise_block, interrupting=interrupting, except_branch=except_branch, reraise_branch=reraise_branch)
+        # This will be the last frame in frames.
+      elif frame.kind == Frame.MODULE:
+        raise_block = frame.blocks['raise_block']
+        block.add_exit(raise_block, interrupting=interrupting, except_branch=except_branch, reraise_branch=reraise_branch)
+        # This will be the last frame in frames.
 
   def new_block(self, node=None, label=None, prunable=True):
     """Create a new block."""
     return self.graph.new_block(node=node, label=label, prunable=prunable)
 
+  def enter_module_frame(self, exit_block, raise_block):
+    # The entire module is in the interior of the frame.
+    # The exit block and raise block are the exits from the frame.
+    self.frames.append(Frame(Frame.MODULE,
+                             exit_block=exit_block,
+                             raise_block=raise_block))
+
   def enter_loop_frame(self, continue_block, break_block):
     # The loop body is the interior of the frame.
     # The continue block (loop condition) and break block (loop's after block)
@@ -807,11 +839,24 @@ def get_current_exception_handling_frames(self):
         # A function frame's raise_block catches any exception that reaches it.
         frames.append(frame)
         return frames
+      if frame.kind == Frame.MODULE:
+        # A module frame's raise_block catches any exception that reaches it.
+        frames.append(frame)
+        return frames
     # There is no frame to fully catch the exception.
-    return None
+    raise ValueError('No frame exists to catch the exception.')
 
   def visit_Module(self, node, current_block):
-    return self.visit_list(node.body, current_block)
+    exit_block = self.new_block(node=node, label='<exit>', prunable=False)
+    raise_block = self.new_block(node=node, label='<raise>', prunable=False)
+    self.enter_module_frame(exit_block, raise_block)
+    end_block = self.visit_list(node.body, current_block)
+    end_block.add_exit(exit_block)
+    self.exit_frame()
+    # Move exit and raise blocks to the end of the block list.
+    self.graph.move_block_to_rear(exit_block)
+    self.graph.move_block_to_rear(raise_block)
+    return end_block
 
   def visit_ClassDef(self, node, current_block):
     """Visit a ClassDef node of the AST.
@@ -821,7 +866,7 @@ def visit_ClassDef(self, node, current_block):
     """
     # TODO(dbieber): Make sure all statements are handled, such as base classes.
     # http://greentreesnakes.readthedocs.io/en/latest/nodes.html#ClassDef
-    # The body is exceuted before the decorators.
+    # The body is executed before the decorators.
     current_block = self.visit_list(node.body, current_block)
     for decorator in node.decorator_list:
       self.add_new_instruction(current_block, decorator)
@@ -895,6 +940,8 @@ def handle_function_definition(self, node, name, args, body):
     fn_block = self.visit_list(body, fn_block)
     fn_block.add_exit(return_block)
     self.exit_frame()
+    self.graph.move_block_to_rear(return_block)
+    self.graph.move_block_to_rear(raise_block)
 
   def handle_argument_defaults(self, node, current_block):
     """Add Instructions for all of a FunctionDef's default values.
@@ -1026,6 +1073,7 @@ def handle_Loop(self, node, loop_instruction, current_block):
     else:
       test_block.add_exit(after_block, branch=False)
 
+    self.graph.move_block_to_rear(after_block)
     return after_block
 
   def visit_Try(self, node, current_block):
@@ -1060,9 +1108,12 @@ def visit_Try(self, node, current_block):
       bare_handler_block = None
 
     if node.finalbody:
+      # TODO(dbieber): Move final_block and all blocks from visiting it
+      # to after try blocks.
       final_block = self.new_block(node=node, label='final_block')
       final_block_end = self.visit_list(node.finalbody, final_block)
-      final_block_end.add_exit(after_block)
+      # "False" indicates the path taken after finally if no error has been raised.
+      final_block_end.add_exit(after_block, reraise_branch=False)
       self.enter_try_finally_frame(final_block, final_block_end)
     else:
       final_block = after_block
@@ -1095,7 +1146,9 @@ def visit_Try(self, node, current_block):
     if bare_handler_block is None and previous_handler_block_end is not None:
       # If no exceptions match, then raise up through the frames.
       # (A bare-except will always match.)
-      self.raise_through_frames(previous_handler_block_end, interrupting=False)
+      # Here "False" indicates the final exception header did not match the raised error.
+      self.raise_through_frames(
+          previous_handler_block_end, interrupting=False, except_branch=False)
 
     if node.orelse:
       else_block = self.visit_list(node.orelse, else_block)
@@ -1104,6 +1157,7 @@ def visit_Try(self, node, current_block):
     if node.finalbody:
       self.exit_frame()  # Exit the try-finally frame.
 
+    self.graph.move_block_to_rear(after_block)
     return after_block
 
   def handle_ExceptHandler(self, handler, handler_block, handler_body_block,
@@ -1134,10 +1188,13 @@ def handle_ExceptHandler(self, handler, handler_block, handler_body_block,
       self.add_new_instruction(handler_block, handler.type)
     # An ExceptHandler header can only have a single Instruction, so there is
     # only one handler_block BasicBlock.
-    handler_block.add_exit(handler_body_block)
+    # Here "True" indicates the exception header matches the raised error.
+    handler_block.add_exit(handler_body_block, except_branch=True)
 
     if previous_handler_block_end is not None:
-      previous_handler_block_end.add_exit(handler_block)
+      # Here "False" indicates the previous exception header did not match the
+      # raised error.
+      previous_handler_block_end.add_exit(handler_block, except_branch=False)
     previous_handler_block_end = handler_block
 
     if handler.name is not None:

setup.py

@@ -30,14 +30,15 @@
     'absl-py',
     'astunparse',
     'gast',
-    'six',
+    'networkx',
     'pygraphviz',
+    'six',
 ]
 
 TEST_DEPENDENCIES = [
 ]
 
-VERSION = '1.0.1'
+VERSION = '1.1.0'
 URL = 'https://github.com/google-research/python-graphs'
 
 setup(