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,26 +458,32 @@ 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)
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
size_t len = arg_bufinfo.len / sz;
// convert byte count to element count
size_t len = arg_bufinfo.len / sz;
// make sure we have enough room to extend
// TODO: alloc policy; at the moment we go conservative
if (self->free < len) {
self->items = m_renew(byte, self->items, (self->len + self->free) * sz, (self->len + len) * sz);
self->free = 0;
} else {
self->free -= len;
}
// make sure we have enough room to extend
// TODO: alloc policy; at the moment we go conservative
if (self->free < len) {
self->items = m_renew(byte, self->items, (self->len + self->free) * sz, (self->len + len) * sz);
self->free = 0;
// extend
mp_seq_copy((byte *)self->items + self->len * sz, arg_bufinfo.buf, len * sz, byte);
self->len += len;
} else {
self->free -= len;
// 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);
}
}
// extend
mp_seq_copy((byte *)self->items + self->len * sz, arg_bufinfo.buf, len * sz, byte);
self->len += len;
return mp_const_none;
}
STATIC MP_DEFINE_CONST_FUN_OBJ_2(array_extend_obj, array_extend);