• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1"""bytecode_helper - support tools for testing correct bytecode generation"""
2
3import unittest
4import dis
5import io
6import opcode
7try:
8    import _testinternalcapi
9except ImportError:
10    _testinternalcapi = None
11
12_UNSPECIFIED = object()
13
14def instructions_with_positions(instrs, co_positions):
15    # Return (instr, positions) pairs from the instrs list and co_positions
16    # iterator. The latter contains items for cache lines and the former
17    # doesn't, so those need to be skipped.
18
19    co_positions = co_positions or iter(())
20    for instr in instrs:
21        yield instr, next(co_positions, ())
22        for _, size, _ in (instr.cache_info or ()):
23            for i in range(size):
24                next(co_positions, ())
25
26class BytecodeTestCase(unittest.TestCase):
27    """Custom assertion methods for inspecting bytecode."""
28
29    def get_disassembly_as_string(self, co):
30        s = io.StringIO()
31        dis.dis(co, file=s)
32        return s.getvalue()
33
34    def assertInBytecode(self, x, opname, argval=_UNSPECIFIED):
35        """Returns instr if opname is found, otherwise throws AssertionError"""
36        self.assertIn(opname, dis.opmap)
37        for instr in dis.get_instructions(x):
38            if instr.opname == opname:
39                if argval is _UNSPECIFIED or instr.argval == argval:
40                    return instr
41        disassembly = self.get_disassembly_as_string(x)
42        if argval is _UNSPECIFIED:
43            msg = '%s not found in bytecode:\n%s' % (opname, disassembly)
44        else:
45            msg = '(%s,%r) not found in bytecode:\n%s'
46            msg = msg % (opname, argval, disassembly)
47        self.fail(msg)
48
49    def assertNotInBytecode(self, x, opname, argval=_UNSPECIFIED):
50        """Throws AssertionError if opname is found"""
51        self.assertIn(opname, dis.opmap)
52        for instr in dis.get_instructions(x):
53            if instr.opname == opname:
54                disassembly = self.get_disassembly_as_string(x)
55                if argval is _UNSPECIFIED:
56                    msg = '%s occurs in bytecode:\n%s' % (opname, disassembly)
57                    self.fail(msg)
58                elif instr.argval == argval:
59                    msg = '(%s,%r) occurs in bytecode:\n%s'
60                    msg = msg % (opname, argval, disassembly)
61                    self.fail(msg)
62
63class CompilationStepTestCase(unittest.TestCase):
64
65    HAS_ARG = set(dis.hasarg)
66    HAS_TARGET = set(dis.hasjrel + dis.hasjabs + dis.hasexc)
67    HAS_ARG_OR_TARGET = HAS_ARG.union(HAS_TARGET)
68
69    class Label:
70        pass
71
72    def assertInstructionsMatch(self, actual_seq, expected):
73        # get an InstructionSequence and an expected list, where each
74        # entry is a label or an instruction tuple. Construct an expcted
75        # instruction sequence and compare with the one given.
76
77        self.assertIsInstance(expected, list)
78        actual = actual_seq.get_instructions()
79        expected = self.seq_from_insts(expected).get_instructions()
80        self.assertEqual(len(actual), len(expected))
81
82        # compare instructions
83        for act, exp in zip(actual, expected):
84            if isinstance(act, int):
85                self.assertEqual(exp, act)
86                continue
87            self.assertIsInstance(exp, tuple)
88            self.assertIsInstance(act, tuple)
89            idx = max([p[0] for p in enumerate(exp) if p[1] != -1])
90            self.assertEqual(exp[:idx], act[:idx])
91
92    def resolveAndRemoveLabels(self, insts):
93        idx = 0
94        res = []
95        for item in insts:
96            assert isinstance(item, (self.Label, tuple))
97            if isinstance(item, self.Label):
98                item.value = idx
99            else:
100                idx += 1
101                res.append(item)
102
103        return res
104
105    def seq_from_insts(self, insts):
106        labels = {item for item in insts if isinstance(item, self.Label)}
107        for i, lbl in enumerate(labels):
108            lbl.value = i
109
110        seq = _testinternalcapi.new_instruction_sequence()
111        for item in insts:
112            if isinstance(item, self.Label):
113                seq.use_label(item.value)
114            else:
115                op = item[0]
116                if isinstance(op, str):
117                    op = opcode.opmap[op]
118                arg, *loc = item[1:]
119                if isinstance(arg, self.Label):
120                    arg = arg.value
121                loc = loc + [-1] * (4 - len(loc))
122                seq.addop(op, arg or 0, *loc)
123        return seq
124
125    def check_instructions(self, insts):
126        for inst in insts:
127            if isinstance(inst, self.Label):
128                continue
129            op, arg, *loc = inst
130            if isinstance(op, str):
131                op = opcode.opmap[op]
132            self.assertEqual(op in opcode.hasarg,
133                             arg is not None,
134                             f"{opcode.opname[op]=} {arg=}")
135            self.assertTrue(all(isinstance(l, int) for l in loc))
136
137
138@unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
139class CodegenTestCase(CompilationStepTestCase):
140
141    def generate_code(self, ast):
142        insts, _ = _testinternalcapi.compiler_codegen(ast, "my_file.py", 0)
143        return insts
144
145
146@unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
147class CfgOptimizationTestCase(CompilationStepTestCase):
148
149    def get_optimized(self, seq, consts, nlocals=0):
150        insts = _testinternalcapi.optimize_cfg(seq, consts, nlocals)
151        return insts, consts
152
153@unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
154class AssemblerTestCase(CompilationStepTestCase):
155
156    def get_code_object(self, filename, insts, metadata):
157        co = _testinternalcapi.assemble_code_object(filename, insts, metadata)
158        return co
159