def map_tuples(typingctx: TypingContext, map_fn: types.Type,
*tuples: types.Type):
# check for accepted types
if isinstance(map_fn, types.Callable) and all(
isinstance(t, types.BaseTuple) for t in tuples):
fn_sigs = [
map_fn.get_call_type(typingctx, tys, {})
for tys in zip(*(t.types for t in tuples))
]
# create the expected type signature
result_type = types.BaseTuple.from_types(
[sig.return_type for sig in fn_sigs])
sig = result_type(map_fn, types.StarArgTuple(tuples))
# defines the custom code generation
def codegen(context: CodegenContext, builder: ir.IRBuilder,
signature: Signature, args: Sequence[ir.Value]):
# llvm IRBuilder code here
map_fn_ty, _ = signature.args
funcs = (context.get_function(map_fn_ty, sig) for sig in fn_sigs)
_, tuples = args
tuples = cgutils.unpack_tuple(builder, tuples)
elems = (func(builder,
[builder.extract_value(t, i) for t in tuples])
for i, func in enumerate(funcs))
return context.make_tuple(builder, signature.return_type, elems)
return sig, codegen
def stararg_handler(index, param, values):
return types.StarArgTuple(values)