1# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
2# For details: https://bitbucket.org/ned/coveragepy/src/default/NOTICE.txt
3
4"""Bytecode manipulation for coverage.py"""
5
6import opcode
7import types
8
9from coverage.backward import byte_to_int
10
11
12class ByteCode(object):
13    """A single bytecode."""
14    def __init__(self):
15        # The offset of this bytecode in the code object.
16        self.offset = -1
17
18        # The opcode, defined in the `opcode` module.
19        self.op = -1
20
21        # The argument, a small integer, whose meaning depends on the opcode.
22        self.arg = -1
23
24        # The offset in the code object of the next bytecode.
25        self.next_offset = -1
26
27        # The offset to jump to.
28        self.jump_to = -1
29
30
31class ByteCodes(object):
32    """Iterator over byte codes in `code`.
33
34    This handles the logic of EXTENDED_ARG byte codes internally.  Those byte
35    codes are not returned by this iterator.
36
37    Returns `ByteCode` objects.
38
39    """
40    def __init__(self, code):
41        self.code = code
42
43    def __getitem__(self, i):
44        return byte_to_int(self.code[i])
45
46    def __iter__(self):
47        offset = 0
48        ext_arg = 0
49        while offset < len(self.code):
50            bc = ByteCode()
51            bc.op = self[offset]
52            bc.offset = offset
53
54            next_offset = offset+1
55            if bc.op >= opcode.HAVE_ARGUMENT:
56                bc.arg = ext_arg + self[offset+1] + 256*self[offset+2]
57                next_offset += 2
58
59                label = -1
60                if bc.op in opcode.hasjrel:
61                    label = next_offset + bc.arg
62                elif bc.op in opcode.hasjabs:
63                    label = bc.arg
64                bc.jump_to = label
65
66            bc.next_offset = offset = next_offset
67            if bc.op == opcode.EXTENDED_ARG:
68                ext_arg = bc.arg * 256*256
69            else:
70                ext_arg = 0
71                yield bc
72
73
74class CodeObjects(object):
75    """Iterate over all the code objects in `code`."""
76    def __init__(self, code):
77        self.stack = [code]
78
79    def __iter__(self):
80        while self.stack:
81            # We're going to return the code object on the stack, but first
82            # push its children for later returning.
83            code = self.stack.pop()
84            for c in code.co_consts:
85                if isinstance(c, types.CodeType):
86                    self.stack.append(c)
87            yield code
88