Skip to content

Commit edbdf71

Browse files
committed
rt_unpack_sequence(): Support generic iterables.
1 parent 48697f1 commit edbdf71

1 file changed

Lines changed: 21 additions & 5 deletions

File tree

py/runtime.c

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -769,26 +769,42 @@ mp_obj_t rt_store_set(mp_obj_t set, mp_obj_t item) {
769769

770770
// unpacked items are stored in reverse order into the array pointed to by items
771771
void rt_unpack_sequence(mp_obj_t seq_in, uint num, mp_obj_t *items) {
772+
uint seq_len;
772773
if (MP_OBJ_IS_TYPE(seq_in, &tuple_type) || MP_OBJ_IS_TYPE(seq_in, &list_type)) {
773-
uint seq_len;
774774
mp_obj_t *seq_items;
775775
if (MP_OBJ_IS_TYPE(seq_in, &tuple_type)) {
776776
mp_obj_tuple_get(seq_in, &seq_len, &seq_items);
777777
} else {
778778
mp_obj_list_get(seq_in, &seq_len, &seq_items);
779779
}
780780
if (seq_len < num) {
781-
nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_ValueError, "need more than %d values to unpack", (void*)(machine_uint_t)seq_len));
781+
goto too_short;
782782
} else if (seq_len > num) {
783-
nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_ValueError, "too many values to unpack (expected %d)", (void*)(machine_uint_t)num));
783+
goto too_long;
784784
}
785785
for (uint i = 0; i < num; i++) {
786786
items[i] = seq_items[num - 1 - i];
787787
}
788788
} else {
789-
// TODO call rt_getiter and extract via rt_iternext
790-
nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_TypeError, "'%s' object is not iterable", mp_obj_get_type_str(seq_in)));
789+
mp_obj_t iterable = rt_getiter(seq_in);
790+
791+
for (seq_len = 0; seq_len < num; seq_len++) {
792+
mp_obj_t el = rt_iternext(iterable);
793+
if (el == mp_const_stop_iteration) {
794+
goto too_short;
795+
}
796+
items[num - 1 - seq_len] = el;
797+
}
798+
if (rt_iternext(iterable) != mp_const_stop_iteration) {
799+
goto too_long;
800+
}
791801
}
802+
return;
803+
804+
too_short:
805+
nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_ValueError, "need more than %d values to unpack", seq_len));
806+
too_long:
807+
nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_ValueError, "too many values to unpack (expected %d)", num));
792808
}
793809

794810
mp_obj_t rt_build_map(int n_args) {

0 commit comments

Comments
 (0)