Skip to content
Go to file
Cannot retrieve contributors at this time
615 lines (467 sloc) 15.6 KB
Python wrapper that connects CPython interpreter to the Numba typed-list.
This is the code that is used when creating typed lists outside of a `@jit`
context and when returning a typed-list from a `@jit` decorated function. It
basically a Python class that has a Numba allocated typed-list under the hood
and uses `@jit` functions to access it. Since it inherits from MutableSequence
it should really quack like the CPython `list`.
from import MutableSequence
from numba.core.types import ListType, TypeRef
from numba.core.imputils import numba_typeref_ctor
from numba.core.dispatcher import Dispatcher
from numba.core import types, config, cgutils
from numba import njit, typeof
from numba.core.extending import (
from numba.typed import listobject
from numba.core.errors import TypingError, LoweringError
from numba.core.typing.templates import Signature
def _make_list(itemty, allocated=DEFAULT_ALLOCATED):
return listobject._as_meminfo(listobject.new_list(itemty,
def _length(l):
return len(l)
def _allocated(l):
return l._allocated()
def _is_mutable(l):
return l._is_mutable()
def _make_mutable(l):
return l._make_mutable()
def _make_immutable(l):
return l._make_immutable()
def _append(l, item):
def _setitem(l, i, item):
l[i] = item
def _getitem(l, i):
return l[i]
def _contains(l, item):
return item in l
def _count(l, item):
return l.count(item)
def _pop(l, i):
return l.pop(i)
def _delitem(l, i):
del l[i]
def _extend(l, iterable):
return l.extend(iterable)
def _insert(l, i, item):
l.insert(i, item)
def _remove(l, item):
def _clear(l):
def _reverse(l):
def _copy(l):
return l.copy()
def _eq(t, o):
return t == o
def _ne(t, o):
return t != o
def _lt(t, o):
return t < o
def _le(t, o):
return t <= o
def _gt(t, o):
return t > o
def _ge(t, o):
return t >= o
def _index(l, item, start, end):
return l.index(item, start, end)
def _sort(l, key, reverse):
return l.sort(key, reverse)
def _from_meminfo_ptr(ptr, listtype):
return List(meminfo=ptr, lsttype=listtype)
class List(MutableSequence):
"""A typed-list usable in Numba compiled functions.
Implements the MutableSequence interface.
_legal_kwargs = ["lsttype", "meminfo", "allocated"]
def __new__(cls,
if config.DISABLE_JIT:
return list.__new__(list)
return object.__new__(cls)
def empty_list(cls, item_type, allocated=DEFAULT_ALLOCATED):
"""Create a new empty List.
item_type: Numba type
type of the list item.
allocated: int
number of items to pre-allocate
if config.DISABLE_JIT:
return list()
return cls(lsttype=ListType(item_type), allocated=allocated)
def __init__(self, *args, **kwargs):
For users, the constructor does not take any parameters.
The keyword arguments are for internal use only.
args: iterable
The iterable to intialize the list from
lsttype : numba.core.types.ListType; keyword-only
Used internally for the list type.
meminfo : MemInfo; keyword-only
Used internally to pass the MemInfo object when boxing.
allocated: int; keyword-only
Used internally to pre-allocate space for items
illegal_kwargs = any((kw not in self._legal_kwargs for kw in kwargs))
if illegal_kwargs or args and kwargs:
raise TypeError("List() takes no keyword arguments")
if kwargs:
self._list_type, self._opaque = self._parse_arg(**kwargs)
self._list_type = None
if args:
if not 0 <= len(args) <= 1:
raise TypeError(
"List() expected at most 1 argument, got {}"
iterable = args[0]
# Special case Numpy scalars or anything that quacks like a
# NumPy Array.
if hasattr(iterable, "ndim") and iterable.ndim == 0:
except TypeError:
raise TypeError("List() argument must be iterable")
for i in args[0]:
def _parse_arg(self, lsttype, meminfo=None, allocated=DEFAULT_ALLOCATED):
if not isinstance(lsttype, ListType):
raise TypeError('*lsttype* must be a ListType')
if meminfo is not None:
opaque = meminfo
opaque = _make_list(lsttype.item_type, allocated=allocated)
return lsttype, opaque
def _numba_type_(self):
if self._list_type is None:
raise TypeError("invalid operation on untyped list")
return self._list_type
def _typed(self):
"""Returns True if the list is typed.
return self._list_type is not None
def _dtype(self):
if not self._typed:
raise RuntimeError("invalid operation on untyped list")
return self._list_type.dtype
def _initialise_list(self, item):
lsttype = types.ListType(typeof(item))
self._list_type, self._opaque = self._parse_arg(lsttype)
def __len__(self):
if not self._typed:
return 0
return _length(self)
def _allocated(self):
if not self._typed:
return _allocated(self)
def _is_mutable(self):
return _is_mutable(self)
def _make_mutable(self):
return _make_mutable(self)
def _make_immutable(self):
return _make_immutable(self)
def __eq__(self, other):
return _eq(self, other)
def __ne__(self, other):
return _ne(self, other)
def __lt__(self, other):
return _lt(self, other)
def __le__(self, other):
return _le(self, other)
def __gt__(self, other):
return _gt(self, other)
def __ge__(self, other):
return _ge(self, other)
def append(self, item):
if not self._typed:
_append(self, item)
def __setitem__(self, i, item):
if not self._typed:
_setitem(self, i, item)
def __getitem__(self, i):
if not self._typed:
raise IndexError
return _getitem(self, i)
def __iter__(self):
for i in range(len(self)):
yield self[i]
def __contains__(self, item):
return _contains(self, item)
def __delitem__(self, i):
_delitem(self, i)
def insert(self, i, item):
if not self._typed:
_insert(self, i, item)
def count(self, item):
return _count(self, item)
def pop(self, i=-1):
return _pop(self, i)
def extend(self, iterable):
# Empty iterable, do nothing
if len(iterable) == 0:
return self
if not self._typed:
# Need to get the first element of the iterable to initialise the
# type of the list. FIXME: this may be a problem if the iterable
# can not be sliced.
return _extend(self, iterable)
def remove(self, item):
return _remove(self, item)
def clear(self):
return _clear(self)
def reverse(self):
return _reverse(self)
def copy(self):
return _copy(self)
def index(self, item, start=None, stop=None):
return _index(self, item, start, stop)
def sort(self, key=None, reverse=False):
"""Sort the list inplace.
See also ``list.sort()``
# If key is not already a dispatcher object, make it so
if callable(key) and not isinstance(key, Dispatcher):
key = njit(key)
return _sort(self, key, reverse)
def __str__(self):
buf = []
for x in self:
return '[{0}]'.format(', '.join(buf))
def __repr__(self):
body = str(self)
prefix = str(self._list_type) if self._typed else "ListType[Undefined]"
return "{prefix}({body})".format(prefix=prefix, body=body)
# XXX: should we have a better way to classmethod
@overload_method(TypeRef, 'empty_list')
def typedlist_empty(cls, item_type, allocated=DEFAULT_ALLOCATED):
if cls.instance_type is not ListType:
def impl(cls, item_type, allocated=DEFAULT_ALLOCATED):
return listobject.new_list(item_type, allocated=allocated)
return impl
def box_lsttype(typ, val, c):
context = c.context
builder = c.builder
# XXX deduplicate
ctor = cgutils.create_struct_proxy(typ)
lstruct = ctor(context, builder, value=val)
# Returns the plain MemInfo
boxed_meminfo =
modname = c.context.insert_const_string(
c.builder.module, 'numba.typed.typedlist',
typedlist_mod = c.pyapi.import_module_noblock(modname)
fmp_fn = c.pyapi.object_getattr_string(typedlist_mod, '_from_meminfo_ptr')
lsttype_obj = c.pyapi.unserialize(c.pyapi.serialize_object(typ))
result_var = builder.alloca(c.pyapi.pyobj), result_var)
with builder.if_then(cgutils.is_not_null(builder, lsttype_obj)):
res = c.pyapi.call_function_objargs(
fmp_fn, (boxed_meminfo, lsttype_obj),
c.pyapi.decref(boxed_meminfo), result_var)
return builder.load(result_var)
def unbox_listtype(typ, val, c):
context = c.context
builder = c.builder
miptr = c.pyapi.object_getattr_string(val, '_opaque')
native = c.unbox(types.MemInfoPointer(types.voidptr), miptr)
mi = native.value
ctor = cgutils.create_struct_proxy(typ)
lstruct = ctor(context, builder)
data_pointer = context.nrt.meminfo_data(builder, mi)
data_pointer = builder.bitcast(
) = builder.load(data_pointer)
lstruct.meminfo = mi
lstobj = lstruct._getvalue()
return NativeValue(lstobj)
# The following contains the logic for the type-inferred constructor
def _guess_dtype(iterable):
"""Guess the correct dtype of the iterable type. """
if not isinstance(iterable, types.IterableType):
raise TypingError(
"List() argument must be iterable")
# Special case for nested NumPy arrays.
elif isinstance(iterable, types.Array) and iterable.ndim > 1:
return iterable.copy(ndim=iterable.ndim - 1, layout='A')
elif hasattr(iterable, "dtype"):
return iterable.dtype
elif hasattr(iterable, "yield_type"):
return iterable.yield_type
elif isinstance(iterable, types.UnicodeType):
return iterable
elif isinstance(iterable, types.DictType):
return iterable.key_type
# This should never happen, since the 'dtype' of any iterable
# should have determined above.
raise TypingError(
"List() argument does not have a suitable dtype")
def typedlist_call(context):
"""Defines typing logic for ``List()`` and ``List(iterable)``.
If no argument is given, the returned typer types a new typed-list with an
undefined item type. If a single argument is given it must be iterable with
a guessable 'dtype'. In this case, the typer types a new typed-list with
the type set to the 'dtype' of the iterable arg.
arg : single iterable (optional)
The single optional argument.
typer : function
A typer suitable to type constructor calls.
The returned typer raises a TypingError in case of unsuitable arguments.
class Typer(object):
def attach_sig(self):
from inspect import signature as mypysig
def mytyper(iterable):
self.pysig = mypysig(mytyper)
def __call__(self, *args, **kwargs):
if kwargs:
raise TypingError(
"List() takes no keyword arguments"
elif args:
if not 0 <= len(args) <= 1:
raise TypingError(
"List() expected at most 1 argument, got {}"
rt = types.ListType(_guess_dtype(args[0]))
return Signature(rt, args, None, pysig=self.pysig)
item_type = types.undefined
return types.ListType(item_type)
return Typer()
def impl_numba_typeref_ctor(cls, *args):
"""Defines lowering for ``List()`` and ``List(iterable)``.
This defines the lowering logic to instantiate either an empty typed-list
or a typed-list initialised with values from a single iterable argument.
cls : TypeRef
Expecting a TypeRef of a precise ListType.
args: tuple
A tuple that contains a single iterable (optional)
impl : function
An implementation suitable for lowering the constructor call.
See also: `redirect_type_ctor` in numba/cpython/
list_ty = cls.instance_type
if not isinstance(list_ty, types.ListType):
return # reject
# Ensure the list is precisely typed.
if not list_ty.is_precise():
msg = "expecting a precise ListType but got {}".format(list_ty)
raise LoweringError(msg)
item_type = types.TypeRef(list_ty.item_type)
if args:
# special case 0d Numpy arrays
if isinstance(args[0], types.Array) and args[0].ndim == 0:
def impl(cls, *args):
# Instatiate an empty list and populate it with the single
# value from the array.
r = List.empty_list(item_type)
return r
def impl(cls, *args):
# Instatiate an empty list and populate it with values from the
# iterable.
r = List.empty_list(item_type)
for i in args[0]:
return r
def impl(cls, *args):
# Simply call .empty_list with the item type from *cls*
return List.empty_list(item_type)
return impl
You can’t perform that action at this time.