From 56beb01724d4f0027babc5d23f016efbde4c4190 Mon Sep 17 00:00:00 2001 From: Paul Sokolovsky Date: Thu, 16 Apr 2015 00:48:36 +0300 Subject: [PATCH] objarray: Support assignment of bytes to bytearray slice. --- py/objarray.c | 32 ++++++++++++++++++-------- tests/basics/bytearray_slice_assign.py | 6 +++++ 2 files changed, 29 insertions(+), 9 deletions(-) diff --git a/py/objarray.c b/py/objarray.c index 8c1caaceaf..cb2b2b6a1a 100644 --- a/py/objarray.c +++ b/py/objarray.c @@ -369,17 +369,31 @@ STATIC mp_obj_t array_subscr(mp_obj_t self_in, mp_obj_t index_in, mp_obj_t value 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; + mp_uint_t src_len; + void *src_items; 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"); + if (MP_OBJ_IS_TYPE(value, &mp_type_array) || MP_OBJ_IS_TYPE(value, &mp_type_bytearray)) { + mp_obj_array_t *src_slice = value; + if (item_sz != mp_binary_get_size('@', src_slice->typecode, NULL)) { + compat_error: + mp_not_implemented("lhs and rhs should be compatible"); + } + src_len = src_slice->len; + src_items = src_slice->items; + } else if (MP_OBJ_IS_TYPE(value, &mp_type_bytes)) { + if (item_sz != 1) { + goto compat_error; + } + mp_buffer_info_t bufinfo; + mp_get_buffer_raise(value, &bufinfo, MP_BUFFER_READ); + src_len = bufinfo.len; + src_items = bufinfo.buf; + } else { + mp_not_implemented("array/bytes required on right side"); } // TODO: check src/dst compat - mp_int_t len_adj = src_slice->len - (slice.stop - slice.start); + mp_int_t len_adj = src_len - (slice.stop - slice.start); if (len_adj > 0) { if (len_adj > o->free) { // TODO: alloc policy; at the moment we go conservative @@ -387,10 +401,10 @@ STATIC mp_obj_t array_subscr(mp_obj_t self_in, mp_obj_t index_in, mp_obj_t value 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); + slice.start, slice.stop, src_items, src_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); + slice.start, slice.stop, src_items, src_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); diff --git a/tests/basics/bytearray_slice_assign.py b/tests/basics/bytearray_slice_assign.py index a5436c815a..3209e24d40 100644 --- a/tests/basics/bytearray_slice_assign.py +++ b/tests/basics/bytearray_slice_assign.py @@ -46,3 +46,9 @@ print(l) b = bytearray(2) b[2:] = bytearray(10) print(b) + + +# Assignment of bytes to array slice +b = bytearray(2) +b[1:1] = b"12345" +print(b)