Skip to content
Permalink
master
Go to file
 
 
Cannot retrieve contributors at this time
67 lines (54 sloc) 1.62 KB
"""
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')
You can’t perform that action at this time.