bytecode_helper.py

"""bytecode_helper - support tools for testing correct bytecode generation"""

importunittest
importdis
importio
importopcode
try:
import_testinternalcapi
exceptImportError:
_testinternalcapi=None

_UNSPECIFIED=object()

definstructions_with_positions(instrs, co_positions):
# Return (instr, positions) pairs from the instrs list and co_positions
# iterator. The latter contains items for cache lines and the former
# doesn't, so those need to be skipped.

co_positions=co_positionsoriter(())
forinstrininstrs:
yieldinstr, next(co_positions, ())
for_, size, _in (instr.cache_infoor ()):
foriinrange(size):
next(co_positions, ())

classBytecodeTestCase(unittest.TestCase):
"""Custom assertion methods for inspecting bytecode."""

defget_disassembly_as_string(self, co):
s=io.StringIO()
dis.dis(co, file=s)
returns.getvalue()

defassertInBytecode(self, x, opname, argval=_UNSPECIFIED):
"""Returns instr if opname is found, otherwise throws AssertionError"""
self.assertIn(opname, dis.opmap)
forinstrindis.get_instructions(x):
ifinstr.opname==opname:
ifargvalis_UNSPECIFIEDorinstr.argval==argval:
returninstr
disassembly=self.get_disassembly_as_string(x)
ifargvalis_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)

defassertNotInBytecode(self, x, opname, argval=_UNSPECIFIED):
"""Throws AssertionError if opname is found"""
self.assertIn(opname, dis.opmap)
forinstrindis.get_instructions(x):
ifinstr.opname==opname:
disassembly=self.get_disassembly_as_string(x)
ifargvalis_UNSPECIFIED:
msg='%s occurs in bytecode:\n%s'% (opname, disassembly)
self.fail(msg)
elifinstr.argval==argval:
msg='(%s,%r) occurs in bytecode:\n%s'
msg=msg% (opname, argval, disassembly)
self.fail(msg)

classCompilationStepTestCase(unittest.TestCase):

HAS_ARG=set(dis.hasarg)
HAS_TARGET=set(dis.hasjrel+dis.hasjabs+dis.hasexc)
HAS_ARG_OR_TARGET=HAS_ARG.union(HAS_TARGET)

classLabel:
pass

defassertInstructionsMatch(self, actual_seq, expected):
# get an InstructionSequence and an expected list, where each
# entry is a label or an instruction tuple. Construct an expcted
# instruction sequence and compare with the one given.

self.assertIsInstance(expected, list)
actual=actual_seq.get_instructions()
expected=self.seq_from_insts(expected).get_instructions()
self.assertEqual(len(actual), len(expected))

# compare instructions
foract, expinzip(actual, expected):
ifisinstance(act, int):
self.assertEqual(exp, act)
continue
self.assertIsInstance(exp, tuple)
self.assertIsInstance(act, tuple)
idx=max([p[0] forpinenumerate(exp) ifp[1] !=-1])
self.assertEqual(exp[:idx], act[:idx])

defresolveAndRemoveLabels(self, insts):
idx=0
res= []
foritemininsts:
assertisinstance(item, (self.Label, tuple))
ifisinstance(item, self.Label):
item.value=idx
else:
idx+=1
res.append(item)

returnres

defseq_from_insts(self, insts):
labels= {itemforitemininstsifisinstance(item, self.Label)}
fori, lblinenumerate(labels):
lbl.value=i

seq=_testinternalcapi.new_instruction_sequence()
foritemininsts:
ifisinstance(item, self.Label):
seq.use_label(item.value)
else:
op=item[0]
ifisinstance(op, str):
op=opcode.opmap[op]
arg, *loc=item[1:]
ifisinstance(arg, self.Label):
arg=arg.value
loc=loc+ [-1] * (4-len(loc))
seq.addop(op, argor0, *loc)
returnseq

defcheck_instructions(self, insts):
forinstininsts:
ifisinstance(inst, self.Label):
continue
op, arg, *loc=inst
ifisinstance(op, str):
op=opcode.opmap[op]
self.assertEqual(opinopcode.hasarg,
argisnotNone,
f"{opcode.opname[op]=}{arg=}")
self.assertTrue(all(isinstance(l, int) forlinloc))


@unittest.skipIf(_testinternalcapiisNone, "requires _testinternalcapi")
classCodegenTestCase(CompilationStepTestCase):

defgenerate_code(self, ast):
insts, _=_testinternalcapi.compiler_codegen(ast, "my_file.py", 0)
returninsts


@unittest.skipIf(_testinternalcapiisNone, "requires _testinternalcapi")
classCfgOptimizationTestCase(CompilationStepTestCase):

defget_optimized(self, seq, consts, nlocals=0):
insts=_testinternalcapi.optimize_cfg(seq, consts, nlocals)
returninsts, consts

@unittest.skipIf(_testinternalcapiisNone, "requires _testinternalcapi")
classAssemblerTestCase(CompilationStepTestCase):

defget_code_object(self, filename, insts, metadata):
co=_testinternalcapi.assemble_code_object(filename, insts, metadata)
returnco