Skip to content
Permalink
master
Switch branches/tags
Go to file
 
 
Cannot retrieve contributors at this time
"""
Implementation of various iterable and iterator types.
"""
from numba.core import types, cgutils
from numba.core.imputils import (
lower_builtin, iternext_impl, call_iternext, call_getiter,
impl_ret_borrowed, impl_ret_new_ref, RefType)
@lower_builtin('getiter', types.IteratorType)
def iterator_getiter(context, builder, sig, args):
[it] = args
return impl_ret_borrowed(context, builder, sig.return_type, it)
#-------------------------------------------------------------------------------
# builtin `enumerate` implementation
@lower_builtin(enumerate, types.IterableType)
@lower_builtin(enumerate, types.IterableType, types.Integer)
def make_enumerate_object(context, builder, sig, args):
assert len(args) == 1 or len(args) == 2 # enumerate(it) or enumerate(it, start)
srcty = sig.args[0]
if len(args) == 1:
src = args[0]
start_val = context.get_constant(types.intp, 0)
elif len(args) == 2:
src = args[0]
start_val = context.cast(builder, args[1], sig.args[1], types.intp)
iterobj = call_getiter(context, builder, srcty, src)
enum = context.make_helper(builder, sig.return_type)
countptr = cgutils.alloca_once(builder, start_val.type)
builder.store(start_val, countptr)
enum.count = countptr
enum.iter = iterobj
res = enum._getvalue()
return impl_ret_new_ref(context, builder, sig.return_type, res)
@lower_builtin('iternext', types.EnumerateType)
@iternext_impl(RefType.NEW)
def iternext_enumerate(context, builder, sig, args, result):
[enumty] = sig.args
[enum] = args
enum = context.make_helper(builder, enumty, value=enum)
count = builder.load(enum.count)
ncount = builder.add(count, context.get_constant(types.intp, 1))
builder.store(ncount, enum.count)
srcres = call_iternext(context, builder, enumty.source_type, enum.iter)
is_valid = srcres.is_valid()
result.set_valid(is_valid)
with builder.if_then(is_valid):
srcval = srcres.yielded_value()
result.yield_(context.make_tuple(builder, enumty.yield_type,
[count, srcval]))
#-------------------------------------------------------------------------------
# builtin `zip` implementation
@lower_builtin(zip, types.VarArg(types.Any))
def make_zip_object(context, builder, sig, args):
zip_type = sig.return_type
assert len(args) == len(zip_type.source_types)
zipobj = context.make_helper(builder, zip_type)
for i, (arg, srcty) in enumerate(zip(args, sig.args)):
zipobj[i] = call_getiter(context, builder, srcty, arg)
res = zipobj._getvalue()
return impl_ret_new_ref(context, builder, sig.return_type, res)
@lower_builtin('iternext', types.ZipType)
@iternext_impl(RefType.NEW)
def iternext_zip(context, builder, sig, args, result):
[zip_type] = sig.args
[zipobj] = args
zipobj = context.make_helper(builder, zip_type, value=zipobj)
if len(zipobj) == 0:
# zip() is an empty iterator
result.set_exhausted()
return
p_ret_tup = cgutils.alloca_once(builder,
context.get_value_type(zip_type.yield_type))
p_is_valid = cgutils.alloca_once_value(builder, value=cgutils.true_bit)
for i, (iterobj, srcty) in enumerate(zip(zipobj, zip_type.source_types)):
is_valid = builder.load(p_is_valid)
# Avoid calling the remaining iternext if a iterator has been exhausted
with builder.if_then(is_valid):
srcres = call_iternext(context, builder, srcty, iterobj)
is_valid = builder.and_(is_valid, srcres.is_valid())
builder.store(is_valid, p_is_valid)
val = srcres.yielded_value()
ptr = cgutils.gep_inbounds(builder, p_ret_tup, 0, i)
builder.store(val, ptr)
is_valid = builder.load(p_is_valid)
result.set_valid(is_valid)
with builder.if_then(is_valid):
result.yield_(builder.load(p_ret_tup))
#-------------------------------------------------------------------------------
# generator implementation
@lower_builtin('iternext', types.Generator)
@iternext_impl(RefType.BORROWED)
def iternext_zip(context, builder, sig, args, result):
genty, = sig.args
gen, = args
impl = context.get_generator_impl(genty)
status, retval = impl(context, builder, sig, args)
context.add_linking_libs(getattr(impl, 'libs', ()))
with cgutils.if_likely(builder, status.is_ok):
result.set_valid(True)
result.yield_(retval)
with cgutils.if_unlikely(builder, status.is_stop_iteration):
result.set_exhausted()
with cgutils.if_unlikely(builder,
builder.and_(status.is_error,
builder.not_(status.is_stop_iteration))):
context.call_conv.return_status_propagate(builder, status)