BasicBlock.__iter__() rejects labels

vstinner · vstinner · commit 7a4b3e07a52b · 2016-03-02T10:08:18.000+01:00
diff --git a/bytecode/cfg.py b/bytecode/cfg.py
@@ -23,11 +23,15 @@ def __iter__(self):
                                  "but %s was found"
                                  % instr.__class__.__name__)
 
-            if index < len(self):
-                if isinstance(instr, Instr) and instr.has_jump():
+            if isinstance(instr, Instr) and instr.has_jump():
+                if index < len(self):
                     raise ValueError("Only the last instruction of a basic "
                                      "block can be a jump")
 
+                if not isinstance(instr.arg, BasicBlock):
+                    raise ValueError("Jump target must a BasicBlock, got %s",
+                                     type(instr.arg).__name__)
+
             yield instr
 
 
@@ -217,7 +221,6 @@ def to_bytecode(self):
                 if isinstance(instr, Instr) and isinstance(instr.arg, BasicBlock):
                     used_blocks.add(id(instr.arg))
 
-
         labels = {}
         jumps = []
         instructions = []
diff --git a/bytecode/tests/test_cfg.py b/bytecode/tests/test_cfg.py
@@ -7,18 +7,27 @@
 
 class BlockTests(unittest.TestCase):
     def test_iter_invalid_types(self):
+        # Labels are not allowed in basic blocks
         block = BasicBlock()
         block.append(Label())
         with self.assertRaises(ValueError):
             list(block)
 
+        # Only one jump allowed and only at the end
         block = BasicBlock()
         block2 = BasicBlock()
         block.extend([Instr('JUMP_ABSOLUTE', block2),
                      Instr('NOP')])
         with self.assertRaises(ValueError):
             list(block)
 
+        # jump target must be a BasicBlock
+        block = BasicBlock()
+        label = Label()
+        block.extend([Instr('JUMP_ABSOLUTE', label)])
+        with self.assertRaises(ValueError):
+            list(block)
+
 
 class BytecodeBlocksTests(TestCase):
     maxDiff = 80 * 100
@@ -136,20 +145,6 @@ def test_to_bytecode(self):
                               Instr('RETURN_VALUE', lineno=3)])
         # FIXME: test other attributes
 
-    def test_to_bytecode_labels(self):
-        blocks = ControlFlowGraph()
-        label = Label()
-        blocks.add_block()
-        blocks[0].extend([Instr('LOAD_NAME', 'test'),
-                          Instr('POP_JUMP_IF_FALSE', label)])
-        blocks[1].append(label)
-
-        with self.assertRaises(ValueError) as cm:
-            blocks.to_bytecode()
-        self.assertEqual(str(cm.exception),
-                         'BasicBlock must only contain SetLineno, Instr and '
-                         'ConcreteInstr objects, but Label was found')
-
     def test_from_bytecode(self):
         bytecode = Bytecode()
         label = Label()
diff --git a/doc/api.rst b/doc/api.rst
@@ -449,13 +449,16 @@ BasicBlock
    It is possible to use concrete instructions (:class:`ConcreteInstr`) in
    blocks, but abstract instructions (:class:`Instr`) are preferred.
 
+   Only the last instruction can have a jump, and the jump target must be a
+   basic block (:class:`BasicBlock`).
+
    Labels (:class:`Label`) must not be used in blocks.
 
    Attributes:
 
    .. attribute:: next_block
 
-      Next block (:class:`BasicBlock`), or ``None``.
+      Next basic block (:class:`BasicBlock`), or ``None``.
 
 
 ControlFlowGraph
@@ -470,8 +473,6 @@ ControlFlowGraph
    entry and no branches out except at the exit. Inherit from
    :class:`BaseBytecode`.
 
-   Jump targets are blocks (:class:`BasicBlock`).
-
    Labels (:class:`Label`) must not be used in blocks.
 
    This class is not designed to emit code, but to analyze and modify existing
