diff options
Diffstat (limited to 'py/objarray.c')
| -rw-r--r-- | py/objarray.c | 43 |
1 files changed, 37 insertions, 6 deletions
diff --git a/py/objarray.c b/py/objarray.c index e2fee1f2ed..e030548bbb 100644 --- a/py/objarray.c +++ b/py/objarray.c @@ -351,17 +351,48 @@ STATIC mp_obj_t array_subscr(mp_obj_t self_in, mp_obj_t index_in, mp_obj_t value if (0) { #if MICROPY_PY_BUILTINS_SLICE } else if (MP_OBJ_IS_TYPE(index_in, &mp_type_slice)) { - if (value != MP_OBJ_SENTINEL) { - // Only getting a slice is suported so far, not assignment - // TODO: confirmed that both bytearray and array.array support - // slice assignment (incl. of different size) - return MP_OBJ_NULL; // op not supported - } mp_bound_slice_t slice; if (!mp_seq_get_fast_slice_indexes(o->len, index_in, &slice)) { nlr_raise(mp_obj_new_exception_msg(&mp_type_NotImplementedError, "only slices with step=1 (aka None) are supported")); } + if (value != MP_OBJ_SENTINEL) { + #if MICROPY_PY_ARRAY_SLICE_ASSIGN + // Assign + if (!MP_OBJ_IS_TYPE(value, &mp_type_array) && !MP_OBJ_IS_TYPE(value, &mp_type_bytearray)) { + mp_not_implemented("array required on right side"); + } + mp_obj_array_t *src_slice = value; + int item_sz = mp_binary_get_size('@', o->typecode, NULL); + if (item_sz != mp_binary_get_size('@', src_slice->typecode, NULL)) { + mp_not_implemented("arrays should be compatible"); + } + + // TODO: check src/dst compat + mp_int_t len_adj = src_slice->len - (slice.stop - slice.start); + if (len_adj > 0) { + if (len_adj > o->free) { + // TODO: alloc policy; at the moment we go conservative + o->items = m_realloc(o->items, (o->len + o->free) * item_sz, (o->len + len_adj) * item_sz); + o->free = 0; + } + mp_seq_replace_slice_grow_inplace(o->items, o->len, + slice.start, slice.stop, src_slice->items, src_slice->len, len_adj, item_sz); + } else { + mp_seq_replace_slice_no_grow(o->items, o->len, + slice.start, slice.stop, src_slice->items, src_slice->len, item_sz); + // Clear "freed" elements at the end of list + // TODO: This is actually only needed for typecode=='O' + mp_seq_clear(o->items, o->len + len_adj, o->len, item_sz); + // TODO: alloc policy after shrinking + } + o->len += len_adj; + return mp_const_none; + #else + return MP_OBJ_NULL; // op not supported + #endif + } + mp_obj_array_t *res; int sz = mp_binary_get_size('@', o->typecode & TYPECODE_MASK, NULL); assert(sz > 0); |