From 32ef3a3517caa79f5d2237dbed8078e7fa79c742 Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Thu, 4 Dec 2014 15:46:14 +0000
Subject: [PATCH] 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).
---
 py/objarray.c                       | 26 ++++++++++++++++++++++----
 py/objstr.c                         |  6 ++++++
 tests/basics/array_construct.py     | 18 ++++++++++++++++++
 tests/basics/bytearray_construct.py | 14 ++++++++++++++
 tests/basics/bytes_construct.py     | 14 ++++++++++++++
 5 files changed, 74 insertions(+), 4 deletions(-)
 create mode 100644 tests/basics/array_construct.py
 create mode 100644 tests/basics/bytearray_construct.py
 create mode 100644 tests/basics/bytes_construct.py

diff --git a/py/objarray.c b/py/objarray.c
index f96d6f423..cb8e03e45 100644
--- a/py/objarray.c
+++ b/py/objarray.c
@@ -126,7 +126,25 @@ STATIC mp_obj_array_t *array_new(char typecode, mp_uint_t n) {
 
 #if MICROPY_PY_BUILTINS_BYTEARRAY || MICROPY_PY_ARRAY
 STATIC mp_obj_t array_construct(char typecode, mp_obj_t initializer) {
-    uint len;
+    // bytearrays can be raw-initialised from anything with the buffer protocol
+    // other arrays can only be raw-initialised from bytes and bytearray objects
+    mp_buffer_info_t bufinfo;
+    if (((MICROPY_PY_BUILTINS_BYTEARRAY
+            && typecode == BYTEARRAY_TYPECODE)
+        || (MICROPY_PY_ARRAY
+            && (MP_OBJ_IS_TYPE(initializer, &mp_type_bytes)
+                || MP_OBJ_IS_TYPE(initializer, &mp_type_bytearray))))
+        && mp_get_buffer(initializer, &bufinfo, MP_BUFFER_READ)) {
+        // construct array from raw bytes
+        // we round-down the len to make it a multiple of sz (CPython raises error)
+        int sz = mp_binary_get_size('@', typecode, NULL);
+        mp_uint_t len = bufinfo.len / sz;
+        mp_obj_array_t *o = array_new(typecode, len);
+        memcpy(o->items, bufinfo.buf, len * sz);
+        return o;
+    }
+
+    mp_uint_t len;
     // Try to create array of exact len if initializer len is known
     mp_obj_t len_in = mp_obj_len_maybe(initializer);
     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
         // 1 arg: make an empty array
         return array_new(*typecode, 0);
     } else {
-        // 2 args: construct the array from the given iterator
+        // 2 args: construct the array from the given object
         return array_construct(*typecode, args[1]);
     }
 }
@@ -179,12 +197,12 @@ STATIC mp_obj_t bytearray_make_new(mp_obj_t type_in, mp_uint_t n_args, mp_uint_t
         return array_new(BYTEARRAY_TYPECODE, 0);
     } else if (MP_OBJ_IS_SMALL_INT(args[0])) {
         // 1 arg, an integer: construct a blank bytearray of that length
-        uint len = MP_OBJ_SMALL_INT_VALUE(args[0]);
+        mp_uint_t len = MP_OBJ_SMALL_INT_VALUE(args[0]);
         mp_obj_array_t *o = array_new(BYTEARRAY_TYPECODE, len);
         memset(o->items, 0, len);
         return o;
     } else {
-        // 1 arg, an iterator: construct the bytearray from that
+        // 1 arg: construct the bytearray from that
         return array_construct(BYTEARRAY_TYPECODE, args[0]);
     }
 }
diff --git a/py/objstr.c b/py/objstr.c
index a8e276025..cfe0ef115 100644
--- a/py/objstr.c
+++ b/py/objstr.c
@@ -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
         return mp_obj_str_builder_end(o);
     }
 
+    // check if argument has the buffer protocol
+    mp_buffer_info_t bufinfo;
+    if (mp_get_buffer(args[0], &bufinfo, MP_BUFFER_READ)) {
+        return mp_obj_new_str_of_type(&mp_type_bytes, bufinfo.buf, bufinfo.len);
+    }
+
     mp_int_t len;
     byte *data;
     vstr_t *vstr = NULL;
diff --git a/tests/basics/array_construct.py b/tests/basics/array_construct.py
new file mode 100644
index 000000000..e050f3003
--- /dev/null
+++ b/tests/basics/array_construct.py
@@ -0,0 +1,18 @@
+# test construction of array.array from different objects
+
+from array import array
+
+# tuple, list
+print(array('b', (1, 2)))
+print(array('h', [1, 2]))
+
+# raw copy from bytes, bytearray
+print(array('h', b'12'))
+print(array('h', bytearray(2)))
+print(array('i', bytearray(4)))
+
+# convert from other arrays
+print(array('H', array('b', [1, 2])))
+print(array('f', array('h', [1, 2])))
+print(array('b', array('I', [1, 2])))
+print(array('d', array('f', [1, 2])))
diff --git a/tests/basics/bytearray_construct.py b/tests/basics/bytearray_construct.py
new file mode 100644
index 000000000..0a7097d55
--- /dev/null
+++ b/tests/basics/bytearray_construct.py
@@ -0,0 +1,14 @@
+# test construction of bytearray from different objects
+
+from array import array
+
+# bytes, tuple, list
+print(bytearray(b'123'))
+print(bytearray((1, 2)))
+print(bytearray([1, 2]))
+
+# arrays
+print(bytearray(array('b', [1, 2])))
+print(bytearray(array('h', [1, 2])))
+print(bytearray(array('I', [1, 2])))
+print(bytearray(array('f', [1, 2.3])))
diff --git a/tests/basics/bytes_construct.py b/tests/basics/bytes_construct.py
new file mode 100644
index 000000000..1eb6d3e48
--- /dev/null
+++ b/tests/basics/bytes_construct.py
@@ -0,0 +1,14 @@
+# test construction of bytes from different objects
+
+from array import array
+
+# tuple, list, bytearray
+print(bytes((1, 2)))
+print(bytes([1, 2]))
+print(bytes(bytearray(4)))
+
+# arrays
+print(bytes(array('b', [1, 2])))
+print(bytes(array('h', [1, 2])))
+print(bytes(array('I', [1, 2])))
+print(bytes(array('f', [1, 2.3])))
-- 
GitLab