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:
committed by
Damien George
parent
c1629dc2ff
commit
0a98f3a911
@@ -113,6 +113,19 @@ static mp_obj_array_t *array_new(char typecode, size_t n) {
|
|||||||
#endif
|
#endif
|
||||||
|
|
||||||
#if MICROPY_PY_BUILTINS_BYTEARRAY || MICROPY_PY_ARRAY
|
#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) {
|
static mp_obj_t array_construct(char typecode, mp_obj_t initializer) {
|
||||||
// bytearrays can be raw-initialised from anything with the buffer protocol
|
// bytearrays can be raw-initialised from anything with the buffer protocol
|
||||||
// other arrays can only be raw-initialised from bytes and bytearray objects
|
// 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_array_t *array = array_new(typecode, len);
|
||||||
|
array_extend_impl(array, initializer, 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);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return MP_OBJ_FROM_PTR(array);
|
return MP_OBJ_FROM_PTR(array);
|
||||||
}
|
}
|
||||||
#endif
|
#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)
|
// 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)) {
|
||||||
|
array_extend_impl(self, arg_in, 0, 0);
|
||||||
|
return mp_const_none;
|
||||||
|
}
|
||||||
|
|
||||||
size_t sz = mp_binary_get_size('@', self->typecode, NULL);
|
size_t sz = mp_binary_get_size('@', self->typecode, NULL);
|
||||||
|
|
||||||
|
|||||||
@@ -14,3 +14,9 @@ print(a1)
|
|||||||
|
|
||||||
a1.extend(array.array('I', [5]))
|
a1.extend(array.array('I', [5]))
|
||||||
print(a1)
|
print(a1)
|
||||||
|
|
||||||
|
a1.extend([6, 7])
|
||||||
|
print(a1)
|
||||||
|
|
||||||
|
a1.extend(i for i in (8, 9))
|
||||||
|
print(a1)
|
||||||
|
|||||||
Reference in New Issue
Block a user