1"""Bytecode manipulation for coverage.py"""
2
3import opcode, sys, types
4
5class ByteCode(object):
6    """A single bytecode."""
7    def __init__(self):
8        # The offset of this bytecode in the code object.
9        self.offset = -1
10
11        # The opcode, defined in the `opcode` module.
12        self.op = -1
13
14        # The argument, a small integer, whose meaning depends on the opcode.
15        self.arg = -1
16
17        # The offset in the code object of the next bytecode.
18        self.next_offset = -1
19
20        # The offset to jump to.
21        self.jump_to = -1
22
23
24class ByteCodes(object):
25    """Iterator over byte codes in `code`.
26
27    Returns `ByteCode` objects.
28
29    """
30    def __init__(self, code):
31        self.code = code
32        self.offset = 0
33
34    if sys.version_info >= (3, 0):
35        def __getitem__(self, i):
36            return self.code[i]
37    else:
38        def __getitem__(self, i):
39            return ord(self.code[i])
40
41    def __iter__(self):
42        return self
43
44    def __next__(self):
45        if self.offset >= len(self.code):
46            raise StopIteration
47
48        bc = ByteCode()
49        bc.op = self[self.offset]
50        bc.offset = self.offset
51
52        next_offset = self.offset+1
53        if bc.op >= opcode.HAVE_ARGUMENT:
54            bc.arg = self[self.offset+1] + 256*self[self.offset+2]
55            next_offset += 2
56
57            label = -1
58            if bc.op in opcode.hasjrel:
59                label = next_offset + bc.arg
60            elif bc.op in opcode.hasjabs:
61                label = bc.arg
62            bc.jump_to = label
63
64        bc.next_offset = self.offset = next_offset
65        return bc
66
67    next = __next__     # Py2k uses an old-style non-dunder name.
68
69
70class CodeObjects(object):
71    """Iterate over all the code objects in `code`."""
72    def __init__(self, code):
73        self.stack = [code]
74
75    def __iter__(self):
76        return self
77
78    def __next__(self):
79        if self.stack:
80            # We're going to return the code object on the stack, but first
81            # push its children for later returning.
82            code = self.stack.pop()
83            for c in code.co_consts:
84                if isinstance(c, types.CodeType):
85                    self.stack.append(c)
86            return code
87
88        raise StopIteration
89
90    next = __next__
91