Skip to content

Commit 56beb01

Browse files
committed
objarray: Support assignment of bytes to bytearray slice.
1 parent 9a18e21 commit 56beb01

2 files changed

Lines changed: 29 additions & 9 deletions

File tree

py/objarray.c

Lines changed: 23 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -369,28 +369,42 @@ STATIC mp_obj_t array_subscr(mp_obj_t self_in, mp_obj_t index_in, mp_obj_t value
369369
if (value != MP_OBJ_SENTINEL) {
370370
#if MICROPY_PY_ARRAY_SLICE_ASSIGN
371371
// Assign
372-
if (!MP_OBJ_IS_TYPE(value, &mp_type_array) && !MP_OBJ_IS_TYPE(value, &mp_type_bytearray)) {
373-
mp_not_implemented("array required on right side");
374-
}
375-
mp_obj_array_t *src_slice = value;
372+
mp_uint_t src_len;
373+
void *src_items;
376374
int item_sz = mp_binary_get_size('@', o->typecode, NULL);
377-
if (item_sz != mp_binary_get_size('@', src_slice->typecode, NULL)) {
378-
mp_not_implemented("arrays should be compatible");
375+
if (MP_OBJ_IS_TYPE(value, &mp_type_array) || MP_OBJ_IS_TYPE(value, &mp_type_bytearray)) {
376+
mp_obj_array_t *src_slice = value;
377+
if (item_sz != mp_binary_get_size('@', src_slice->typecode, NULL)) {
378+
compat_error:
379+
mp_not_implemented("lhs and rhs should be compatible");
380+
}
381+
src_len = src_slice->len;
382+
src_items = src_slice->items;
383+
} else if (MP_OBJ_IS_TYPE(value, &mp_type_bytes)) {
384+
if (item_sz != 1) {
385+
goto compat_error;
386+
}
387+
mp_buffer_info_t bufinfo;
388+
mp_get_buffer_raise(value, &bufinfo, MP_BUFFER_READ);
389+
src_len = bufinfo.len;
390+
src_items = bufinfo.buf;
391+
} else {
392+
mp_not_implemented("array/bytes required on right side");
379393
}
380394

381395
// TODO: check src/dst compat
382-
mp_int_t len_adj = src_slice->len - (slice.stop - slice.start);
396+
mp_int_t len_adj = src_len - (slice.stop - slice.start);
383397
if (len_adj > 0) {
384398
if (len_adj > o->free) {
385399
// TODO: alloc policy; at the moment we go conservative
386400
o->items = m_renew(byte, o->items, (o->len + o->free) * item_sz, (o->len + len_adj) * item_sz);
387401
o->free = 0;
388402
}
389403
mp_seq_replace_slice_grow_inplace(o->items, o->len,
390-
slice.start, slice.stop, src_slice->items, src_slice->len, len_adj, item_sz);
404+
slice.start, slice.stop, src_items, src_len, len_adj, item_sz);
391405
} else {
392406
mp_seq_replace_slice_no_grow(o->items, o->len,
393-
slice.start, slice.stop, src_slice->items, src_slice->len, item_sz);
407+
slice.start, slice.stop, src_items, src_len, item_sz);
394408
// Clear "freed" elements at the end of list
395409
// TODO: This is actually only needed for typecode=='O'
396410
mp_seq_clear(o->items, o->len + len_adj, o->len, item_sz);

tests/basics/bytearray_slice_assign.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,3 +46,9 @@
4646
b = bytearray(2)
4747
b[2:] = bytearray(10)
4848
print(b)
49+
50+
51+
# Assignment of bytes to array slice
52+
b = bytearray(2)
53+
b[1:1] = b"12345"
54+
print(b)

0 commit comments

Comments
 (0)