add support for array.extend(iterable)

This commit is contained in:
Gregory Neverov 2023-02-06 15:07:26 -08:00
parent 5f43a63a70
commit abc5eb30d7

View File

@ -458,8 +458,7 @@ STATIC mp_obj_t array_extend(mp_obj_t self_in, mp_obj_t arg_in) {
// allow to extend by anything that has the buffer protocol (extension to CPython) // allow to extend by anything that has the buffer protocol (extension to CPython)
mp_buffer_info_t arg_bufinfo; mp_buffer_info_t arg_bufinfo;
mp_get_buffer_raise(arg_in, &arg_bufinfo, MP_BUFFER_READ); if (mp_get_buffer(arg_in, &arg_bufinfo, MP_BUFFER_READ)) {
size_t sz = mp_binary_get_size('@', self->typecode, NULL); size_t sz = mp_binary_get_size('@', self->typecode, NULL);
// convert byte count to element count // convert byte count to element count
@ -477,7 +476,14 @@ STATIC mp_obj_t array_extend(mp_obj_t self_in, mp_obj_t arg_in) {
// extend // extend
mp_seq_copy((byte *)self->items + self->len * sz, arg_bufinfo.buf, len * sz, byte); mp_seq_copy((byte *)self->items + self->len * sz, arg_bufinfo.buf, len * sz, byte);
self->len += len; self->len += len;
} else {
// Otherwise argument must be an iterable of items to append
mp_obj_t iterable = mp_getiter(arg_in, NULL);
mp_obj_t item;
while ((item = mp_iternext(iterable)) != MP_OBJ_STOP_ITERATION) {
array_append(self_in, item);
}
}
return mp_const_none; return mp_const_none;
} }
STATIC MP_DEFINE_CONST_FUN_OBJ_2(array_extend_obj, array_extend); STATIC MP_DEFINE_CONST_FUN_OBJ_2(array_extend_obj, array_extend);