52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
|
import unittest
|
||
|
|
|
||
|
|
import dis
|
||
|
|
import struct
|
||
|
|
|
||
|
|
from numba import jit
|
||
|
|
from numba.core import utils
|
||
|
|
from numba.tests.support import TestCase
|
||
|
|
|
||
|
|
|
||
|
|
class TestExtendedArg(TestCase):
|
||
|
|
"""
|
||
|
|
Test support for the EXTENDED_ARG opcode.
|
||
|
|
"""
|
||
|
|
bytecode_len = 0xff
|
||
|
|
|
||
|
|
def get_extended_arg_load_const(self):
|
||
|
|
"""
|
||
|
|
Get a function with a EXTENDED_ARG opcode before a LOAD_CONST opcode.
|
||
|
|
"""
|
||
|
|
def f():
|
||
|
|
x = 5
|
||
|
|
return x
|
||
|
|
|
||
|
|
b = bytearray(f.__code__.co_code)
|
||
|
|
consts = f.__code__.co_consts
|
||
|
|
bytecode_format = "<BB"
|
||
|
|
consts = consts + (None,) * self.bytecode_len + (42,)
|
||
|
|
if utils.PYVERSION >= (3, 11):
|
||
|
|
# Python 3.11 has a RESUME op code at the start of a function, need
|
||
|
|
# to inject the EXTENDED_ARG after this to influence the LOAD_CONST
|
||
|
|
offset = 2 # 2 byte op code
|
||
|
|
else:
|
||
|
|
offset = 0
|
||
|
|
|
||
|
|
packed_extend_arg = struct.pack(bytecode_format, dis.EXTENDED_ARG, 1)
|
||
|
|
b[:] = b[:offset] + packed_extend_arg + b[offset:]
|
||
|
|
f.__code__ = f.__code__.replace(co_code=bytes(b), co_consts=consts)
|
||
|
|
return f
|
||
|
|
|
||
|
|
def test_extended_arg_load_const(self):
|
||
|
|
pyfunc = self.get_extended_arg_load_const()
|
||
|
|
# make sure that the pyfunc contains the expected modification
|
||
|
|
self.assertGreater(len(pyfunc.__code__.co_consts), self.bytecode_len)
|
||
|
|
self.assertPreciseEqual(pyfunc(), 42)
|
||
|
|
cfunc = jit(nopython=True)(pyfunc)
|
||
|
|
self.assertPreciseEqual(cfunc(), 42)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|