124 lines
3.8 KiB
Python
124 lines
3.8 KiB
Python
|
import ctypes
|
||
|
|
import unittest
|
||
|
|
from numba.core import types
|
||
|
|
from numba.core.extending import intrinsic
|
||
|
|
from numba import jit, njit
|
||
|
|
from numba.tests.support import captured_stdout
|
||
|
|
|
||
|
|
|
||
|
|
@intrinsic
|
||
|
|
def _pyapi_bytes_as_string(typingctx, csrc, size):
|
||
|
|
sig = types.voidptr(csrc, size) # cstring == void*
|
||
|
|
|
||
|
|
def codegen(context, builder, sig, args):
|
||
|
|
[csrc, size] = args
|
||
|
|
api = context.get_python_api(builder)
|
||
|
|
b = api.bytes_from_string_and_size(csrc, size)
|
||
|
|
return api.bytes_as_string(b)
|
||
|
|
return sig, codegen
|
||
|
|
|
||
|
|
|
||
|
|
def PyBytes_AsString(uni):
|
||
|
|
# test_PyBytes_AsString will call this function with a unicode type.
|
||
|
|
# We then use the underlying buffer to create a PyBytes object and call the
|
||
|
|
# PyBytes_AsString function with PyBytes object as argument
|
||
|
|
return _pyapi_bytes_as_string(uni._data, uni._length)
|
||
|
|
|
||
|
|
|
||
|
|
@intrinsic
|
||
|
|
def _pyapi_bytes_as_string_and_size(typingctx, csrc, size):
|
||
|
|
# return a tuple containing the c-string and size
|
||
|
|
retty = types.Tuple.from_types((csrc, size))
|
||
|
|
sig = retty(csrc, size)
|
||
|
|
|
||
|
|
def codegen(context, builder, sig, args):
|
||
|
|
[csrc, size] = args
|
||
|
|
pyapi = context.get_python_api(builder)
|
||
|
|
b = pyapi.bytes_from_string_and_size(csrc, size)
|
||
|
|
p_cstr = builder.alloca(pyapi.cstring)
|
||
|
|
p_size = builder.alloca(pyapi.py_ssize_t)
|
||
|
|
pyapi.bytes_as_string_and_size(b, p_cstr, p_size)
|
||
|
|
|
||
|
|
cstr = builder.load(p_cstr)
|
||
|
|
size = builder.load(p_size)
|
||
|
|
tup = context.make_tuple(builder, sig.return_type, (cstr, size))
|
||
|
|
return tup
|
||
|
|
return sig, codegen
|
||
|
|
|
||
|
|
|
||
|
|
def PyBytes_AsStringAndSize(uni):
|
||
|
|
return _pyapi_bytes_as_string_and_size(uni._data, uni._length)
|
||
|
|
|
||
|
|
|
||
|
|
class TestPythonAPI(unittest.TestCase):
|
||
|
|
|
||
|
|
def test_PyBytes_AsString(self):
|
||
|
|
cfunc = jit(nopython=True)(PyBytes_AsString)
|
||
|
|
cstr = cfunc('hello') # returns a cstring
|
||
|
|
|
||
|
|
fn = ctypes.pythonapi.PyBytes_FromString
|
||
|
|
fn.argtypes = [ctypes.c_void_p]
|
||
|
|
fn.restype = ctypes.py_object
|
||
|
|
obj = fn(cstr)
|
||
|
|
|
||
|
|
# Use the cstring created from bytes_as_string to create a python
|
||
|
|
# bytes object
|
||
|
|
self.assertEqual(obj, b'hello')
|
||
|
|
|
||
|
|
def test_PyBytes_AsStringAndSize(self):
|
||
|
|
cfunc = jit(nopython=True)(PyBytes_AsStringAndSize)
|
||
|
|
tup = cfunc('hello\x00world') # returns a tuple: cstring and its size
|
||
|
|
|
||
|
|
fn = ctypes.pythonapi.PyBytes_FromStringAndSize
|
||
|
|
fn.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
|
||
|
|
fn.restype = ctypes.py_object
|
||
|
|
obj = fn(tup[0], tup[1])
|
||
|
|
|
||
|
|
# Use the cstring created from bytes_from_string_and_size to create
|
||
|
|
# a python bytes object
|
||
|
|
self.assertEqual(obj, b'hello\x00world')
|
||
|
|
|
||
|
|
|
||
|
|
class PythonAPIEmptyArgs(unittest.TestCase):
|
||
|
|
def test_empty_args(self):
|
||
|
|
def callme(**kwargs):
|
||
|
|
print("callme", kwargs)
|
||
|
|
|
||
|
|
@intrinsic
|
||
|
|
def py_call(tyctx):
|
||
|
|
def codegen(context, builder, sig, args):
|
||
|
|
pyapi = context.get_python_api(builder)
|
||
|
|
gil = pyapi.gil_ensure()
|
||
|
|
|
||
|
|
num = pyapi.long_from_longlong(
|
||
|
|
context.get_constant(types.intp, 0xCAFE)
|
||
|
|
)
|
||
|
|
kwds = pyapi.dict_pack({"key": num}.items())
|
||
|
|
fn_print = pyapi.unserialize(pyapi.serialize_object(callme))
|
||
|
|
# segfault: https://github.com/numba/numba/issues/5871
|
||
|
|
res = pyapi.call(fn_print, None, kwds)
|
||
|
|
|
||
|
|
pyapi.decref(res)
|
||
|
|
pyapi.decref(fn_print)
|
||
|
|
pyapi.decref(kwds)
|
||
|
|
pyapi.decref(num)
|
||
|
|
|
||
|
|
pyapi.gil_release(gil)
|
||
|
|
return res
|
||
|
|
|
||
|
|
return types.none(), codegen
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo():
|
||
|
|
py_call()
|
||
|
|
|
||
|
|
with captured_stdout() as out:
|
||
|
|
foo()
|
||
|
|
d = {"key": 0xCAFE}
|
||
|
|
expected = f"callme {d}\n"
|
||
|
|
self.assertEqual(out.getvalue(), expected)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|