diff --git a/Lib/test/bytecode_helper.py b/Lib/test/bytecode_helper.py new file mode 100644 index 00000000000..c4943cda00d --- /dev/null +++ b/Lib/test/bytecode_helper.py @@ -0,0 +1,72 @@ +"""bytecode_helper - support tools for testing correct bytecode generation""" + +import unittest +import dis +import io + +_UNSPECIFIED = object() + +class BytecodeTestCase(unittest.TestCase): + """Custom assertion methods for inspecting bytecode.""" + + def get_disassembly_as_string(self, co): + s = io.StringIO() + dis.dis(co, file=s) + return s.getvalue() + + def assertInstructionMatches(self, instr, expected, *, line_offset=0): + # Deliberately test opname first, since that gives a more + # meaningful error message than testing opcode + self.assertEqual(instr.opname, expected.opname) + self.assertEqual(instr.opcode, expected.opcode) + self.assertEqual(instr.arg, expected.arg) + self.assertEqual(instr.argval, expected.argval) + self.assertEqual(instr.argrepr, expected.argrepr) + self.assertEqual(instr.offset, expected.offset) + if expected.starts_line is None: + self.assertIsNone(instr.starts_line) + else: + self.assertEqual(instr.starts_line, + expected.starts_line + line_offset) + self.assertEqual(instr.is_jump_target, expected.is_jump_target) + + + def assertBytecodeExactlyMatches(self, x, expected, *, line_offset=0): + """Throws AssertionError if any discrepancy is found in bytecode + + *x* is the object to be introspected + *expected* is a list of dis.Instruction objects + + Set *line_offset* as appropriate to adjust for the location of the + object to be disassembled within the test file. If the expected list + assumes the first line is line 1, then an appropriate offset would be + ``1 - f.__code__.co_firstlineno``. + """ + actual = dis.get_instructions(x, line_offset=line_offset) + self.assertEqual(list(actual), expected) + + def assertInBytecode(self, x, opname, argval=_UNSPECIFIED): + """Returns instr if op is found, otherwise throws AssertionError""" + for instr in dis.get_instructions(x): + if instr.opname == opname: + if argval is _UNSPECIFIED or instr.argval == argval: + return instr + disassembly = self.get_disassembly_as_string(x) + if argval is _UNSPECIFIED: + msg = '%s not found in bytecode:\n%s' % (opname, disassembly) + else: + msg = '(%s,%r) not found in bytecode:\n%s' + msg = msg % (opname, argval, disassembly) + self.fail(msg) + + def assertNotInBytecode(self, x, opname, argval=_UNSPECIFIED): + """Throws AssertionError if op is found""" + for instr in dis.get_instructions(x): + if instr.opname == opname: + disassembly = self.get_disassembly_as_string(co) + if opargval is _UNSPECIFIED: + msg = '%s occurs in bytecode:\n%s' % (opname, disassembly) + elif instr.argval == argval: + msg = '(%s,%r) occurs in bytecode:\n%s' + msg = msg % (opname, argval, disassembly) + self.fail(msg)