py: Generalise and reduce code size of array +, += and .extend().
By using the buffer protocol for these array operations, we now allow addition of memoryview objects, and objects with "incompatible" typecodes (in this case it just adds bytes naively). This is an extension to CPython which seems sensible. It also reduces the code size.
This commit is contained in:
parent
d8c2b2a1c4
commit
c7ca01ad96
@ -229,22 +229,20 @@ STATIC mp_obj_t array_binary_op(mp_uint_t op, mp_obj_t lhs_in, mp_obj_t rhs_in)
|
|||||||
mp_obj_array_t *lhs = lhs_in;
|
mp_obj_array_t *lhs = lhs_in;
|
||||||
switch (op) {
|
switch (op) {
|
||||||
case MP_BINARY_OP_ADD: {
|
case MP_BINARY_OP_ADD: {
|
||||||
#if MICROPY_PY_BUILTINS_MEMORYVIEW
|
// allow to add anything that has the buffer protocol (extension to CPython)
|
||||||
if (lhs->base.type == &mp_type_memoryview) {
|
mp_buffer_info_t lhs_bufinfo;
|
||||||
return MP_OBJ_NULL; // op not supported
|
mp_buffer_info_t rhs_bufinfo;
|
||||||
}
|
array_get_buffer(lhs_in, &lhs_bufinfo, MP_BUFFER_READ);
|
||||||
#endif
|
mp_get_buffer_raise(rhs_in, &rhs_bufinfo, MP_BUFFER_READ);
|
||||||
// if we get here then lhs is not a memoryview, so we don't need to use (& TYPECODE_MASK)
|
|
||||||
if (mp_obj_get_type(rhs_in) != lhs->base.type) {
|
int sz = mp_binary_get_size('@', lhs_bufinfo.typecode, NULL);
|
||||||
return MP_OBJ_NULL; // op not supported
|
|
||||||
}
|
// convert byte count to element count (in case rhs is not multiple of sz)
|
||||||
mp_obj_array_t *rhs = rhs_in;
|
mp_uint_t rhs_len = rhs_bufinfo.len / sz;
|
||||||
if (lhs->typecode != rhs->typecode) {
|
|
||||||
return MP_OBJ_NULL; // op not supported
|
// note: lhs->len is element count of lhs, lhs_bufinfo.len is byte count
|
||||||
}
|
mp_obj_array_t *res = array_new(lhs_bufinfo.typecode, lhs->len + rhs_len);
|
||||||
int sz = mp_binary_get_size('@', lhs->typecode, NULL);
|
mp_seq_cat((byte*)res->items, lhs_bufinfo.buf, lhs_bufinfo.len, rhs_bufinfo.buf, rhs_len * sz, byte);
|
||||||
mp_obj_array_t *res = array_new(lhs->typecode, lhs->len + rhs->len);
|
|
||||||
mp_seq_cat((byte*)res->items, lhs->items, lhs->len * sz, rhs->items, rhs->len * sz, byte);
|
|
||||||
return res;
|
return res;
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -297,32 +295,27 @@ STATIC mp_obj_t array_extend(mp_obj_t self_in, mp_obj_t arg_in) {
|
|||||||
assert(MP_OBJ_IS_TYPE(self_in, &mp_type_array) || MP_OBJ_IS_TYPE(self_in, &mp_type_bytearray));
|
assert(MP_OBJ_IS_TYPE(self_in, &mp_type_array) || MP_OBJ_IS_TYPE(self_in, &mp_type_bytearray));
|
||||||
mp_obj_array_t *self = self_in;
|
mp_obj_array_t *self = self_in;
|
||||||
|
|
||||||
// check for compatible types (array & array, or bytearray & bytearray)
|
// allow to extend by anything that has the buffer protocol (extension to CPython)
|
||||||
if (mp_obj_get_type(arg_in) != self->base.type) {
|
mp_buffer_info_t arg_bufinfo;
|
||||||
type_error:
|
mp_get_buffer_raise(arg_in, &arg_bufinfo, MP_BUFFER_READ);
|
||||||
nlr_raise(mp_obj_new_exception_msg(&mp_type_TypeError,
|
|
||||||
"incompatible type for array operation"));
|
|
||||||
}
|
|
||||||
|
|
||||||
// check for compatible typecode
|
|
||||||
mp_obj_array_t *arg = arg_in;
|
|
||||||
if (self->typecode != arg->typecode) {
|
|
||||||
goto type_error;
|
|
||||||
}
|
|
||||||
|
|
||||||
int sz = mp_binary_get_size('@', self->typecode, NULL);
|
int sz = mp_binary_get_size('@', self->typecode, NULL);
|
||||||
|
|
||||||
|
// convert byte count to element count
|
||||||
|
mp_uint_t len = arg_bufinfo.len / sz;
|
||||||
|
|
||||||
// make sure we have enough room to extend
|
// make sure we have enough room to extend
|
||||||
if (self->free < arg->len) {
|
// TODO: alloc policy; at the moment we go conservative
|
||||||
// TODO: alloc policy; at the moment we go conservative
|
if (self->free < len) {
|
||||||
self->items = m_realloc(self->items, (self->len + self->free) * sz, (self->len + arg->len) * sz);
|
self->items = m_realloc(self->items, (self->len + self->free) * sz, (self->len + len) * sz);
|
||||||
self->free += arg->len;
|
self->free = 0;
|
||||||
|
} else {
|
||||||
|
self->free -= len;
|
||||||
}
|
}
|
||||||
|
|
||||||
// extend
|
// extend
|
||||||
mp_seq_copy((byte*)self->items + self->len * sz, arg->items, arg->len * sz, byte);
|
mp_seq_copy((byte*)self->items + self->len * sz, arg_bufinfo.buf, len * sz, byte);
|
||||||
self->len += arg->len;
|
self->len += len;
|
||||||
self->free -= arg->len;
|
|
||||||
|
|
||||||
return mp_const_none;
|
return mp_const_none;
|
||||||
}
|
}
|
||||||
|
Loading…
Reference in New Issue
Block a user