Skip to content

Commit c698d26

Browse files
committed
list: Add extend() methods and += operator.
1 parent eae1644 commit c698d26

2 files changed

Lines changed: 34 additions & 0 deletions

File tree

py/objlist.c

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ typedef struct _mp_obj_list_t {
2121

2222
static mp_obj_t mp_obj_new_list_iterator(mp_obj_list_t *list, int cur);
2323
static mp_obj_list_t *list_new(uint n);
24+
static mp_obj_t list_extend(mp_obj_t self_in, mp_obj_t arg_in);
2425

2526
/******************************************************************************/
2627
/* list */
@@ -81,6 +82,14 @@ static mp_obj_t list_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
8182
memcpy(s->items + o->len, p->items, sizeof(mp_obj_t) * p->len);
8283
return s;
8384
}
85+
case RT_BINARY_OP_INPLACE_ADD:
86+
{
87+
if (!MP_OBJ_IS_TYPE(rhs, &list_type)) {
88+
return NULL;
89+
}
90+
list_extend(lhs, rhs);
91+
return o;
92+
}
8493
case RT_BINARY_OP_MULTIPLY:
8594
{
8695
if (!MP_OBJ_IS_SMALL_INT(rhs)) {
@@ -117,6 +126,23 @@ mp_obj_t mp_obj_list_append(mp_obj_t self_in, mp_obj_t arg) {
117126
return mp_const_none; // return None, as per CPython
118127
}
119128

129+
static mp_obj_t list_extend(mp_obj_t self_in, mp_obj_t arg_in) {
130+
assert(MP_OBJ_IS_TYPE(self_in, &list_type));
131+
assert(MP_OBJ_IS_TYPE(arg_in, &list_type));
132+
mp_obj_list_t *self = self_in;
133+
mp_obj_list_t *arg = arg_in;
134+
135+
if (self->len + arg->len > self->alloc) {
136+
// TODO: use alloc policy for "4"
137+
self->items = m_renew(mp_obj_t, self->items, self->alloc, self->len + arg->len + 4);
138+
self->alloc = self->len + arg->len + 4;
139+
}
140+
141+
memcpy(self->items + self->len, arg->items, sizeof(mp_obj_t) * arg->len);
142+
self->len += arg->len;
143+
return mp_const_none; // return None, as per CPython
144+
}
145+
120146
static mp_obj_t list_pop(int n_args, const mp_obj_t *args) {
121147
assert(1 <= n_args && n_args <= 2);
122148
assert(MP_OBJ_IS_TYPE(args[0], &list_type));
@@ -281,6 +307,7 @@ static mp_obj_t list_reverse(mp_obj_t self_in) {
281307
}
282308

283309
static MP_DEFINE_CONST_FUN_OBJ_2(list_append_obj, mp_obj_list_append);
310+
static MP_DEFINE_CONST_FUN_OBJ_2(list_extend_obj, list_extend);
284311
static MP_DEFINE_CONST_FUN_OBJ_1(list_clear_obj, list_clear);
285312
static MP_DEFINE_CONST_FUN_OBJ_1(list_copy_obj, list_copy);
286313
static MP_DEFINE_CONST_FUN_OBJ_2(list_count_obj, list_count);
@@ -296,6 +323,7 @@ static const mp_method_t list_type_methods[] = {
296323
{ "clear", &list_clear_obj },
297324
{ "copy", &list_copy_obj },
298325
{ "count", &list_count_obj },
326+
{ "extend", &list_extend_obj },
299327
{ "index", &list_index_obj },
300328
{ "insert", &list_insert_obj },
301329
{ "pop", &list_pop_obj },

tests/basics/tests/list1.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,3 +10,9 @@
1010
f = x.append
1111
f(4)
1212
print(x)
13+
14+
x.extend([100, 200])
15+
print(x)
16+
17+
x += [2, 1]
18+
print(x)

0 commit comments

Comments
 (0)