objarray: Support assignment of bytes to bytearray slice.

1 files changed, 23 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);