diff options
| -rw-r--r-- | py/objlist.c | 40 | ||||
| -rw-r--r-- | tests/basics/list1.py | 2 |
2 files changed, 26 insertions, 16 deletions
diff --git a/py/objlist.c b/py/objlist.c index 531e4b85bb..7c29867332 100644 --- a/py/objlist.c +++ b/py/objlist.c @@ -39,6 +39,15 @@ STATIC void list_print(void (*print)(void *env, const char *fmt, ...), void *env print(env, "]"); } +STATIC mp_obj_t list_extend_from_iter(mp_obj_t list, mp_obj_t iterable) { + mp_obj_t iter = mp_getiter(iterable); + mp_obj_t item; + while ((item = mp_iternext(iter)) != MP_OBJ_NULL) { + mp_obj_list_append(list, item); + } + return list; +} + STATIC mp_obj_t list_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp_obj_t *args) { // TODO check n_kw == 0 @@ -50,13 +59,9 @@ STATIC mp_obj_t list_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp case 1: { // make list from iterable - mp_obj_t iterable = mp_getiter(args[0]); + // TODO: optimize list/tuple mp_obj_t list = mp_obj_new_list(0, NULL); - mp_obj_t item; - while ((item = mp_iternext(iterable)) != MP_OBJ_NULL) { - mp_obj_list_append(list, item); - } - return list; + return list_extend_from_iter(list, args[0]); } default: @@ -173,18 +178,21 @@ mp_obj_t mp_obj_list_append(mp_obj_t self_in, mp_obj_t arg) { STATIC mp_obj_t list_extend(mp_obj_t self_in, mp_obj_t arg_in) { assert(MP_OBJ_IS_TYPE(self_in, &mp_type_list)); - assert(MP_OBJ_IS_TYPE(arg_in, &mp_type_list)); - mp_obj_list_t *self = self_in; - mp_obj_list_t *arg = arg_in; + if (MP_OBJ_IS_TYPE(arg_in, &mp_type_list)) { + mp_obj_list_t *self = self_in; + mp_obj_list_t *arg = arg_in; + + if (self->len + arg->len > self->alloc) { + // TODO: use alloc policy for "4" + self->items = m_renew(mp_obj_t, self->items, self->alloc, self->len + arg->len + 4); + self->alloc = self->len + arg->len + 4; + } - if (self->len + arg->len > self->alloc) { - // TODO: use alloc policy for "4" - self->items = m_renew(mp_obj_t, self->items, self->alloc, self->len + arg->len + 4); - self->alloc = self->len + arg->len + 4; + memcpy(self->items + self->len, arg->items, sizeof(mp_obj_t) * arg->len); + self->len += arg->len; + } else { + list_extend_from_iter(self_in, arg_in); } - - memcpy(self->items + self->len, arg->items, sizeof(mp_obj_t) * arg->len); - self->len += arg->len; return mp_const_none; // return None, as per CPython } diff --git a/tests/basics/list1.py b/tests/basics/list1.py index 122f61d493..a0c8afc4f8 100644 --- a/tests/basics/list1.py +++ b/tests/basics/list1.py @@ -13,6 +13,8 @@ print(x) x.extend([100, 200]) print(x) +x.extend(range(3)) +print(x) x += [2, 1] print(x) |