| Nick Coghlan | 37c7465 | 2013-05-07 08:28:21 +1000 | [diff] [blame] | 1 | """bytecode_helper - support tools for testing correct bytecode generation""" |
| 2 | |
| 3 | import unittest |
| 4 | import dis |
| 5 | import io |
| 6 | |
| 7 | _UNSPECIFIED = object() |
| 8 | |
| 9 | class BytecodeTestCase(unittest.TestCase): |
| 10 | """Custom assertion methods for inspecting bytecode.""" |
| 11 | |
| 12 | def get_disassembly_as_string(self, co): |
| 13 | s = io.StringIO() |
| 14 | dis.dis(co, file=s) |
| 15 | return s.getvalue() |
| 16 | |
| 17 | def assertInstructionMatches(self, instr, expected, *, line_offset=0): |
| 18 | # Deliberately test opname first, since that gives a more |
| 19 | # meaningful error message than testing opcode |
| 20 | self.assertEqual(instr.opname, expected.opname) |
| 21 | self.assertEqual(instr.opcode, expected.opcode) |
| 22 | self.assertEqual(instr.arg, expected.arg) |
| 23 | self.assertEqual(instr.argval, expected.argval) |
| 24 | self.assertEqual(instr.argrepr, expected.argrepr) |
| 25 | self.assertEqual(instr.offset, expected.offset) |
| 26 | if expected.starts_line is None: |
| 27 | self.assertIsNone(instr.starts_line) |
| 28 | else: |
| 29 | self.assertEqual(instr.starts_line, |
| 30 | expected.starts_line + line_offset) |
| 31 | self.assertEqual(instr.is_jump_target, expected.is_jump_target) |
| 32 | |
| 33 | |
| 34 | def assertBytecodeExactlyMatches(self, x, expected, *, line_offset=0): |
| 35 | """Throws AssertionError if any discrepancy is found in bytecode |
| 36 | |
| 37 | *x* is the object to be introspected |
| 38 | *expected* is a list of dis.Instruction objects |
| 39 | |
| 40 | Set *line_offset* as appropriate to adjust for the location of the |
| 41 | object to be disassembled within the test file. If the expected list |
| 42 | assumes the first line is line 1, then an appropriate offset would be |
| 43 | ``1 - f.__code__.co_firstlineno``. |
| 44 | """ |
| 45 | actual = dis.get_instructions(x, line_offset=line_offset) |
| 46 | self.assertEqual(list(actual), expected) |
| 47 | |
| 48 | def assertInBytecode(self, x, opname, argval=_UNSPECIFIED): |
| 49 | """Returns instr if op is found, otherwise throws AssertionError""" |
| 50 | for instr in dis.get_instructions(x): |
| 51 | if instr.opname == opname: |
| 52 | if argval is _UNSPECIFIED or instr.argval == argval: |
| 53 | return instr |
| 54 | disassembly = self.get_disassembly_as_string(x) |
| 55 | if argval is _UNSPECIFIED: |
| 56 | msg = '%s not found in bytecode:\n%s' % (opname, disassembly) |
| 57 | else: |
| 58 | msg = '(%s,%r) not found in bytecode:\n%s' |
| 59 | msg = msg % (opname, argval, disassembly) |
| 60 | self.fail(msg) |
| 61 | |
| 62 | def assertNotInBytecode(self, x, opname, argval=_UNSPECIFIED): |
| 63 | """Throws AssertionError if op is found""" |
| 64 | for instr in dis.get_instructions(x): |
| 65 | if instr.opname == opname: |
| 66 | disassembly = self.get_disassembly_as_string(co) |
| 67 | if opargval is _UNSPECIFIED: |
| 68 | msg = '%s occurs in bytecode:\n%s' % (opname, disassembly) |
| 69 | elif instr.argval == argval: |
| 70 | msg = '(%s,%r) occurs in bytecode:\n%s' |
| 71 | msg = msg % (opname, argval, disassembly) |
| 72 | self.fail(msg) |