Skip to content
Permalink
 
 
Cannot retrieve contributors at this time
264 lines (200 sloc) 6.58 KB
from collections import namedtuple
from functools import singledispatch
import ctypes
import enum
import numpy as np
from numba.core import types, utils, errors
from numba.np import numpy_support
# terminal color markup
_termcolor = errors.termcolor()
class Purpose(enum.Enum):
# Value being typed is used as an argument
argument = 1
# Value being typed is used as a constant
constant = 2
_TypeofContext = namedtuple("_TypeofContext", ("purpose",))
def typeof(val, purpose=Purpose.argument):
"""
Get the Numba type of a Python value for the given purpose.
"""
# Note the behaviour for Purpose.argument must match _typeof.c.
c = _TypeofContext(purpose)
ty = typeof_impl(val, c)
if ty is None:
msg = _termcolor.errmsg(
f"Cannot determine Numba type of {type(val)}")
raise ValueError(msg)
return ty
@singledispatch
def typeof_impl(val, c):
"""
Generic typeof() implementation.
"""
tp = _typeof_buffer(val, c)
if tp is not None:
return tp
# cffi is handled here as it does not expose a public base class
# for exported functions or CompiledFFI instances.
from numba.core.typing import cffi_utils
if cffi_utils.SUPPORTED:
if cffi_utils.is_cffi_func(val):
return cffi_utils.make_function_type(val)
if cffi_utils.is_ffi_instance(val):
return types.ffi
return getattr(val, "_numba_type_", None)
def _typeof_buffer(val, c):
from numba.core.typing import bufproto
try:
m = memoryview(val)
except TypeError:
return
# Object has the buffer protocol
try:
dtype = bufproto.decode_pep3118_format(m.format, m.itemsize)
except ValueError:
return
type_class = bufproto.get_type_class(type(val))
layout = bufproto.infer_layout(m)
return type_class(dtype, m.ndim, layout=layout,
readonly=m.readonly)
@typeof_impl.register(ctypes._CFuncPtr)
def _typeof_ctypes_function(val, c):
from .ctypes_utils import is_ctypes_funcptr, make_function_type
if is_ctypes_funcptr(val):
return make_function_type(val)
@typeof_impl.register(type)
def _typeof_type(val, c):
"""
Type various specific Python types.
"""
if issubclass(val, BaseException):
return types.ExceptionClass(val)
if issubclass(val, tuple) and hasattr(val, "_asdict"):
return types.NamedTupleClass(val)
if issubclass(val, np.generic):
return types.NumberClass(numpy_support.from_dtype(val))
from numba.typed import Dict
if issubclass(val, Dict):
return types.TypeRef(types.DictType)
from numba.typed import List
if issubclass(val, List):
return types.TypeRef(types.ListType)
@typeof_impl.register(bool)
def _typeof_bool(val, c):
return types.boolean
@typeof_impl.register(float)
def _typeof_float(val, c):
return types.float64
@typeof_impl.register(complex)
def _typeof_complex(val, c):
return types.complex128
@typeof_impl.register(int)
def _typeof_int(val, c):
# As in _typeof.c
nbits = utils.bit_length(val)
if nbits < 32:
typ = types.intp
elif nbits < 64:
typ = types.int64
elif nbits == 64 and val >= 0:
typ = types.uint64
else:
raise ValueError("Int value is too large: %s" % val)
return typ
@typeof_impl.register(np.generic)
def _typeof_numpy_scalar(val, c):
try:
return numpy_support.map_arrayscalar_type(val)
except NotImplementedError:
pass
@typeof_impl.register(str)
def _typeof_str(val, c):
return types.string
@typeof_impl.register(type((lambda a: a).__code__))
def _typeof_code(val, c):
return types.code_type
@typeof_impl.register(type(None))
def _typeof_none(val, c):
return types.none
@typeof_impl.register(type(Ellipsis))
def _typeof_ellipsis(val, c):
return types.ellipsis
@typeof_impl.register(tuple)
def _typeof_tuple(val, c):
tys = [typeof_impl(v, c) for v in val]
if any(ty is None for ty in tys):
return
return types.BaseTuple.from_types(tys, type(val))
@typeof_impl.register(list)
def _typeof_list(val, c):
if len(val) == 0:
raise ValueError("Cannot type empty list")
ty = typeof_impl(val[0], c)
if ty is None:
raise ValueError(
f"Cannot type list element type {type(val[0])}")
return types.List(ty, reflected=True)
@typeof_impl.register(set)
def _typeof_set(val, c):
if len(val) == 0:
raise ValueError("Cannot type empty set")
item = next(iter(val))
ty = typeof_impl(item, c)
if ty is None:
raise ValueError(
f"Cannot type set element type {type(item)}")
return types.Set(ty, reflected=True)
@typeof_impl.register(slice)
def _typeof_slice(val, c):
return types.slice2_type if val.step in (None, 1) else types.slice3_type
@typeof_impl.register(enum.Enum)
@typeof_impl.register(enum.IntEnum)
def _typeof_enum(val, c):
clsty = typeof_impl(type(val), c)
return clsty.member_type
@typeof_impl.register(enum.EnumMeta)
def _typeof_enum_class(val, c):
cls = val
members = list(cls.__members__.values())
if len(members) == 0:
raise ValueError("Cannot type enum with no members")
dtypes = {typeof_impl(mem.value, c) for mem in members}
if len(dtypes) > 1:
raise ValueError("Cannot type heterogeneous enum: "
"got value types %s"
% ", ".join(sorted(str(ty) for ty in dtypes)))
if issubclass(val, enum.IntEnum):
typecls = types.IntEnumClass
else:
typecls = types.EnumClass
return typecls(cls, dtypes.pop())
@typeof_impl.register(np.dtype)
def _typeof_dtype(val, c):
tp = numpy_support.from_dtype(val)
return types.DType(tp)
@typeof_impl.register(np.ndarray)
def _typeof_ndarray(val, c):
try:
dtype = numpy_support.from_dtype(val.dtype)
except NotImplementedError:
raise ValueError("Unsupported array dtype: %s" % (val.dtype,))
layout = numpy_support.map_layout(val)
readonly = not val.flags.writeable
return types.Array(dtype, val.ndim, layout, readonly=readonly)
@typeof_impl.register(types.NumberClass)
def _typeof_number_class(val, c):
return val
@typeof_impl.register(types.Literal)
def _typeof_literal(val, c):
return val
@typeof_impl.register(types.TypeRef)
def _typeof_typeref(val, c):
return val
@typeof_impl.register(types.Type)
def _typeof_nb_type(val, c):
if isinstance(val, types.BaseFunction):
return val
elif isinstance(val, (types.Number, types.Boolean)):
return types.NumberClass(val)
else:
return types.TypeRef(val)
You can’t perform that action at this time.