diff options
| -rw-r--r-- | py/runtime.c | 26 |
1 files changed, 21 insertions, 5 deletions
diff --git a/py/runtime.c b/py/runtime.c index 6f6e3c903e..a3970fe3dc 100644 --- a/py/runtime.c +++ b/py/runtime.c @@ -769,8 +769,8 @@ mp_obj_t rt_store_set(mp_obj_t set, mp_obj_t item) { // unpacked items are stored in reverse order into the array pointed to by items void rt_unpack_sequence(mp_obj_t seq_in, uint num, mp_obj_t *items) { + uint seq_len; if (MP_OBJ_IS_TYPE(seq_in, &tuple_type) || MP_OBJ_IS_TYPE(seq_in, &list_type)) { - uint seq_len; mp_obj_t *seq_items; if (MP_OBJ_IS_TYPE(seq_in, &tuple_type)) { mp_obj_tuple_get(seq_in, &seq_len, &seq_items); @@ -778,17 +778,33 @@ void rt_unpack_sequence(mp_obj_t seq_in, uint num, mp_obj_t *items) { mp_obj_list_get(seq_in, &seq_len, &seq_items); } if (seq_len < num) { - nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_ValueError, "need more than %d values to unpack", (void*)(machine_uint_t)seq_len)); + goto too_short; } else if (seq_len > num) { - nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_ValueError, "too many values to unpack (expected %d)", (void*)(machine_uint_t)num)); + goto too_long; } for (uint i = 0; i < num; i++) { items[i] = seq_items[num - 1 - i]; } } else { - // TODO call rt_getiter and extract via rt_iternext - nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_TypeError, "'%s' object is not iterable", mp_obj_get_type_str(seq_in))); + mp_obj_t iterable = rt_getiter(seq_in); + + for (seq_len = 0; seq_len < num; seq_len++) { + mp_obj_t el = rt_iternext(iterable); + if (el == mp_const_stop_iteration) { + goto too_short; + } + items[num - 1 - seq_len] = el; + } + if (rt_iternext(iterable) != mp_const_stop_iteration) { + goto too_long; + } } + return; + +too_short: + nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_ValueError, "need more than %d values to unpack", seq_len)); +too_long: + nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_ValueError, "too many values to unpack (expected %d)", num)); } mp_obj_t rt_build_map(int n_args) { |