Correct handling of breaks and continues.

3 files changed, 75 insertions, 22 deletions

diff --git a/coverage/bytecode.py b/coverage/bytecode.py
index 82d327c..62a19ba 100644
--- a/coverage/bytecode.py
+++ b/coverage/bytecode.py
@@ -7,7 +7,7 @@ class ByteCode(object):
     def __init__(self):
         self.offset = -1
         self.op = -1
-        self.oparg = -1
+        self.arg = -1
         self.next_offset = -1
         self.jump_to = -1
 
@@ -42,14 +42,14 @@ class ByteCodes(object):
         
         next_offset = self.offset+1
         if bc.op >= opcode.HAVE_ARGUMENT:
-            bc.oparg = self[self.offset+1] + 256*self[self.offset+2]
+            bc.arg = self[self.offset+1] + 256*self[self.offset+2]
             next_offset += 2
             
             label = -1
             if bc.op in opcode.hasjrel:
-                label = next_offset + bc.oparg
+                label = next_offset + bc.arg
             elif bc.op in opcode.hasjabs:
-                label = bc.oparg
+                label = bc.arg
             bc.jump_to = label
             
         bc.next_offset = self.offset = next_offset
diff --git a/coverage/parser.py b/coverage/parser.py
index 8884e41..99656f1 100644
--- a/coverage/parser.py
+++ b/coverage/parser.py
@@ -171,8 +171,23 @@ class CodeParser:
         arcs = self.byte_parser._all_arcs()
         return sorted(arcs)
 
+# Opcodes that guide the ByteParser.
+
+def _opcode_set(*names):
+    return set([opcode.opmap[name] for name in names])
+
+OPS_CODE_END = _opcode_set('RETURN_VALUE')
+OPS_CHUNK_END = _opcode_set(
+    'JUMP_ABSOLUTE', 'JUMP_FORWARD', 'RETURN_VALUE', 'RAISE_VARARGS',
+    'BREAK_LOOP', 'CONTINUE_LOOP',
+    )
+OPS_PUSH_BLOCK = _opcode_set('SETUP_LOOP', 'SETUP_EXCEPT', 'SETUP_FINALLY')
+OPS_POP_BLOCK = _opcode_set('POP_BLOCK')
+OPS_BREAK = _opcode_set('BREAK_LOOP')
+
 
 class ByteParser:
+    """Parse byte codes to understand the structure of code."""
 
     def __init__(self, code=None, text=None, filename=None):
 
@@ -273,18 +288,11 @@ class ByteParser:
                 last_line = l
         return last_line
 
-    _code_enders = set([opcode.opmap[name] for name in ['RETURN_VALUE']])
-    _chunk_enders = set([opcode.opmap[name] for name in [
-        'JUMP_ABSOLUTE', 'JUMP_FORWARD', 'BREAK_LOOP', 'CONTINUE_LOOP',
-        'RAISE_VARARGS'
-        ]])
-    _chunk_enders |= _code_enders
-    
     def _split_into_chunks(self):
         class Chunk(object):
             """A sequence of bytecodes with exits to other bytecodes.
             
-            An exit of -1 means the chunk can leave the code block (return).
+            An exit of -1 means the chunk can leave the code (return).
             
             """
             def __init__(self, byte, line=0):
@@ -296,10 +304,14 @@ class ByteParser:
             def __repr__(self):
                 return "<%d:%d(%d) %r>" % (self.byte, self.line, self.length, list(self.exits))
 
+        # The list of chunks so far, and the one we're working on.
         chunks = []
         chunk = None
         bytes_lines_map = dict(self._bytes_lines())
-
+        
+        # The block stack: just store the "handler" value. 
+        block_stack = []
+        
         for bc in ByteCodes(self.code.co_code):
             # Maybe have to start a new block
             if bc.offset in bytes_lines_map:
@@ -311,20 +323,26 @@ class ByteParser:
             if not chunk:
                 chunk = Chunk(bc.offset)
                 chunks.append(chunk)
-                
+
+            # Look at the opcode                
             if bc.jump_to >= 0:
                 chunk.exits.add(bc.jump_to)
             
-            if bc.op in self._code_enders:
+            if bc.op in OPS_CODE_END:
                 chunk.exits.add(-1)
-
-            if bc.op in self._chunk_enders:
+            elif bc.op in OPS_PUSH_BLOCK:
+                block_stack.append(bc.jump_to)
+            elif bc.op in OPS_POP_BLOCK:
+                block_stack.pop()
+            elif bc.op in OPS_CHUNK_END:
+                if bc.op in OPS_BREAK:
+                    chunk.exits.add(block_stack[-1])
                 chunk = None
-                
+
         if chunks:
             chunks[-1].length = bc.next_offset - chunks[-1].byte
-        for i in range(len(chunks)-1):
-            chunks[i].length = chunks[i+1].byte - chunks[i].byte
+            for i in range(len(chunks)-1):
+                chunks[i].length = chunks[i+1].byte - chunks[i].byte
 
         return chunks
 
@@ -383,7 +401,8 @@ class ByteParser:
             if chunk.line:
                 for ex in chunk.exits:
                     for exit_line in byte_lines[ex]:
-                        arcs.add((chunk.line, exit_line))
+                        if chunk.line != exit_line:
+                            arcs.add((chunk.line, exit_line))
         for line in byte_lines[0]:
             arcs.add((-1, line))
         
diff --git a/test/test_arcs.py b/test/test_arcs.py
index 48e1b28..d26f30b 100644
--- a/test/test_arcs.py
+++ b/test/test_arcs.py
@@ -9,7 +9,7 @@ sys.path.insert(0, os.path.split(__file__)[0]) # Force relative import for Py3k
 from coveragetest import CoverageTest
 
 
-class ArcTest(CoverageTest):
+class SimpleArcTest(CoverageTest):
     """Tests for Coverage.py's arc measurement."""
 
     def test_simple_sequence(self):
@@ -76,6 +76,9 @@ class ArcTest(CoverageTest):
             """,
             arcz=".1 12 25 14 45 5.", arcz_missing="12 25")
         
+class LoopArcTest(CoverageTest):
+    """Arc-measuring tests involving loops."""
+    
     def test_loop(self):
         self.check_coverage("""\
             for i in range(10):
@@ -91,6 +94,15 @@ class ArcTest(CoverageTest):
             """,
             arcz=".1 12 23 32 24 4.", arcz_missing="23 32")
 
+    def test_nested_loop(self):
+        self.check_coverage("""\
+            for i in range(3):
+                for j in range(3):
+                    a = i + j
+            assert a == 4
+            """,
+            arcz=".1 12 23 32 21 14 4.", arcz_missing="")
+
     def test_break(self):
         self.check_coverage("""\
             for i in range(10):
@@ -101,6 +113,28 @@ class ArcTest(CoverageTest):
             """,
             arcz=".1 12 23 35 15 41 5.", arcz_missing="15 41")
 
+    def test_continue(self):
+        self.check_coverage("""\
+            for i in range(10):
+                a = i
+                continue    # 3
+                a = 99
+            assert a == 9   # 5
+            """,
+            arcz=".1 12 23 31 15 41 5.", arcz_missing="41")
+
+    def test_nested_breaks(self):
+        self.check_coverage("""\
+            for i in range(3):
+                for j in range(3):
+                    a = i + j
+                    break               # 4
+                if i == 2:
+                    break
+            assert a == 2 and i == 2    # 7
+            """,
+            arcz=".1 12 23 34 45 25 56 51 67 17 7.", arcz_missing="17 25")
+        
     def xest_xx(self):
         self.check_coverage("""\
             """,