Skip to content
Permalink
master
Switch branches/tags
Go to file
 
 
Cannot retrieve contributors at this time
"""
Typing support for the buffer protocol (PEP 3118).
"""
import array
from numba.core import types
_pep3118_int_types = set('bBhHiIlLqQnN')
_pep3118_scalar_map = {
'f': types.float32,
'd': types.float64,
'Zf': types.complex64,
'Zd': types.complex128,
}
_type_map = {
bytearray: types.ByteArray,
array.array: types.PyArray,
}
_type_map[memoryview] = types.MemoryView
_type_map[bytes] = types.Bytes
def decode_pep3118_format(fmt, itemsize):
"""
Return the Numba type for an item with format string *fmt* and size
*itemsize* (in bytes).
"""
# XXX reuse _dtype_from_pep3118() from np.core._internal?
if fmt in _pep3118_int_types:
# Determine int width and signedness
name = 'int%d' % (itemsize * 8,)
if fmt.isupper():
name = 'u' + name
return types.Integer(name)
try:
# For the hard-coded types above, consider "=" the same as "@"
# (the default). This is because Numpy sometimes adds "="
# in front of the PEP 3118 format string.
return _pep3118_scalar_map[fmt.lstrip('=')]
except KeyError:
raise ValueError("unsupported PEP 3118 format %r" % (fmt,))
def get_type_class(typ):
"""
Get the Numba type class for buffer-compatible Python *typ*.
"""
try:
# Look up special case.
return _type_map[typ]
except KeyError:
# Fall back on generic one.
return types.Buffer
def infer_layout(val):
"""
Infer layout of the given memoryview *val*.
"""
return ('C' if val.c_contiguous else
'F' if val.f_contiguous else
'A')