Skip to content

Commit 32ef3a3

Browse files
dpgeorgedpgeorge
committed
py: Allow bytes/bytearray/array to be init'd by buffer protocol objects.
Behaviour of array initialisation is subtly different for bytes, bytearray and array.array when argument has buffer protocol. This patch gets us CPython conformant (except we allow initialisation of array.array by buffer with length not a multiple of typecode).
1 parent 3a5352b commit 32ef3a3

5 files changed

Lines changed: 74 additions & 4 deletions

File tree

‎py/objarray.c‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,25 @@ STATIC mp_obj_array_t *array_new(char typecode, mp_uint_t n) {
126126

127127
#if MICROPY_PY_BUILTINS_BYTEARRAY || MICROPY_PY_ARRAY
128128
STATIC mp_obj_t array_construct(char typecode, mp_obj_t initializer) {
129-
uint len;
129+
// bytearrays can be raw-initialised from anything with the buffer protocol
130+
// other arrays can only be raw-initialised from bytes and bytearray objects
131+
mp_buffer_info_t bufinfo;
132+
if (((MICROPY_PY_BUILTINS_BYTEARRAY
133+
&& typecode == BYTEARRAY_TYPECODE)
134+
|| (MICROPY_PY_ARRAY
135+
&& (MP_OBJ_IS_TYPE(initializer, &mp_type_bytes)
136+
|| MP_OBJ_IS_TYPE(initializer, &mp_type_bytearray))))
137+
&& mp_get_buffer(initializer, &bufinfo, MP_BUFFER_READ)) {
138+
// construct array from raw bytes
139+
// we round-down the len to make it a multiple of sz (CPython raises error)
140+
int sz = mp_binary_get_size('@', typecode, NULL);
141+
mp_uint_t len = bufinfo.len / sz;
142+
mp_obj_array_t *o = array_new(typecode, len);
143+
memcpy(o->items, bufinfo.buf, len * sz);
144+
return o;
145+
}
146+
147+
mp_uint_t len;
130148
// Try to create array of exact len if initializer len is known
131149
mp_obj_t len_in = mp_obj_len_maybe(initializer);
132150
if (len_in == MP_OBJ_NULL) {
@@ -164,7 +182,7 @@ STATIC mp_obj_t array_make_new(mp_obj_t type_in, mp_uint_t n_args, mp_uint_t n_k
164182
// 1 arg: make an empty array
165183
return array_new(*typecode, 0);
166184
} else {
167-
// 2 args: construct the array from the given iterator
185+
// 2 args: construct the array from the given object
168186
return array_construct(*typecode, args[1]);
169187
}
170188
}
@@ -179,12 +197,12 @@ STATIC mp_obj_t bytearray_make_new(mp_obj_t type_in, mp_uint_t n_args, mp_uint_t
179197
return array_new(BYTEARRAY_TYPECODE, 0);
180198
} else if (MP_OBJ_IS_SMALL_INT(args[0])) {
181199
// 1 arg, an integer: construct a blank bytearray of that length
182-
uint len = MP_OBJ_SMALL_INT_VALUE(args[0]);
200+
mp_uint_t len = MP_OBJ_SMALL_INT_VALUE(args[0]);
183201
mp_obj_array_t *o = array_new(BYTEARRAY_TYPECODE, len);
184202
memset(o->items, 0, len);
185203
return o;
186204
} else {
187-
// 1 arg, an iterator: construct the bytearray from that
205+
// 1 arg: construct the bytearray from that
188206
return array_construct(BYTEARRAY_TYPECODE, args[0]);
189207
}
190208
}

‎py/objstr.c‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,12 @@ STATIC mp_obj_t bytes_make_new(mp_obj_t type_in, mp_uint_t n_args, mp_uint_t n_k
212212
return mp_obj_str_builder_end(o);
213213
}
214214

215+
// check if argument has the buffer protocol
216+
mp_buffer_info_t bufinfo;
217+
if (mp_get_buffer(args[0], &bufinfo, MP_BUFFER_READ)) {
218+
return mp_obj_new_str_of_type(&mp_type_bytes, bufinfo.buf, bufinfo.len);
219+
}
220+
215221
mp_int_t len;
216222
byte *data;
217223
vstr_t *vstr = NULL;

‎tests/basics/array_construct.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
# test construction of array.array from different objects
2+
3+
from array import array
4+
5+
# tuple, list
6+
print(array('b', (1, 2)))
7+
print(array('h', [1, 2]))
8+
9+
# raw copy from bytes, bytearray
10+
print(array('h', b'12'))
11+
print(array('h', bytearray(2)))
12+
print(array('i', bytearray(4)))
13+
14+
# convert from other arrays
15+
print(array('H', array('b', [1, 2])))
16+
print(array('f', array('h', [1, 2])))
17+
print(array('b', array('I', [1, 2])))
18+
print(array('d', array('f', [1, 2])))

‎tests/basics/bytearray_construct.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
# test construction of bytearray from different objects
2+
3+
from array import array
4+
5+
# bytes, tuple, list
6+
print(bytearray(b'123'))
7+
print(bytearray((1, 2)))
8+
print(bytearray([1, 2]))
9+
10+
# arrays
11+
print(bytearray(array('b', [1, 2])))
12+
print(bytearray(array('h', [1, 2])))
13+
print(bytearray(array('I', [1, 2])))
14+
print(bytearray(array('f', [1, 2.3])))

‎tests/basics/bytes_construct.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
# test construction of bytes from different objects
2+
3+
from array import array
4+
5+
# tuple, list, bytearray
6+
print(bytes((1, 2)))
7+
print(bytes([1, 2]))
8+
print(bytes(bytearray(4)))
9+
10+
# arrays
11+
print(bytes(array('b', [1, 2])))
12+
print(bytes(array('h', [1, 2])))
13+
print(bytes(array('I', [1, 2])))
14+
print(bytes(array('f', [1, 2.3])))

0 commit comments

Comments
 (0)