py/objarray: Allow extending array with any iterable.

As suggested by @dpgeorge, factor out part of array_construct to allow it
to be used for construction & extension.

Note that extending with a known-length list (or tuple) goes through the
slow path of calling array_extend once per element.

Fixes issue #7408.

Signed-off-by: Jeff Epler <jepler@gmail.com>
This commit is contained in:
Jeff Epler
2025-06-01 19:19:51 +02:00
committed by Damien George
parent c1629dc2ff
commit 0a98f3a911
2 changed files with 24 additions and 13 deletions

View File

@@ -113,6 +113,19 @@ static mp_obj_array_t *array_new(char typecode, size_t n) {
#endif
#if MICROPY_PY_BUILTINS_BYTEARRAY || MICROPY_PY_ARRAY
static void array_extend_impl(mp_obj_array_t *array, mp_obj_t arg, char typecode, size_t len) {
mp_obj_t iterable = mp_getiter(arg, NULL);
mp_obj_t item;
size_t i = 0;
while ((item = mp_iternext(iterable)) != MP_OBJ_STOP_ITERATION) {
if (len == 0) {
array_append(MP_OBJ_FROM_PTR(array), item);
} else {
mp_binary_set_val_array(typecode, array->items, i++, item);
}
}
}
static mp_obj_t array_construct(char typecode, mp_obj_t initializer) {
// bytearrays can be raw-initialised from anything with the buffer protocol
// other arrays can only be raw-initialised from bytes and bytearray objects
@@ -142,18 +155,7 @@ static mp_obj_t array_construct(char typecode, mp_obj_t initializer) {
}
mp_obj_array_t *array = array_new(typecode, len);
mp_obj_t iterable = mp_getiter(initializer, NULL);
mp_obj_t item;
size_t i = 0;
while ((item = mp_iternext(iterable)) != MP_OBJ_STOP_ITERATION) {
if (len == 0) {
array_append(MP_OBJ_FROM_PTR(array), item);
} else {
mp_binary_set_val_array(typecode, array->items, i++, item);
}
}
array_extend_impl(array, initializer, typecode, len);
return MP_OBJ_FROM_PTR(array);
}
#endif
@@ -413,7 +415,10 @@ 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)) {
array_extend_impl(self, arg_in, 0, 0);
return mp_const_none;
}
size_t sz = mp_binary_get_size('@', self->typecode, NULL);

View File

@@ -14,3 +14,9 @@ print(a1)
a1.extend(array.array('I', [5]))
print(a1)
a1.extend([6, 7])
print(a1)
a1.extend(i for i in (8, 9))
print(a1)