py: Implement memoryview slice assignment.

dxxb · dpgeorge · commit 32aba4083032 · 2015-07-20T15:53:22.000+01:00
Adds ability to do "memcpy" with memoryview objects, such as:
m1[0:3] = m2[2:5].
diff --git a/py/objarray.c b/py/objarray.c
@@ -373,12 +373,12 @@ STATIC mp_obj_t array_subscr(mp_obj_t self_in, mp_obj_t index_in, mp_obj_t value
                 // Assign
                 mp_uint_t src_len;
                 void *src_items;
-                size_t item_sz = mp_binary_get_size('@', o->typecode, NULL);
+                size_t item_sz = mp_binary_get_size('@', o->typecode & TYPECODE_MASK, NULL);
                 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)) {
+                    if (item_sz != mp_binary_get_size('@', src_slice->typecode & TYPECODE_MASK, NULL)) {
                     compat_error:
-                        mp_not_implemented("lhs and rhs should be compatible");
+                        nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "lhs and rhs should be compatible"));
                     }
                     src_len = src_slice->len;
                     src_items = src_slice->items;
@@ -390,13 +390,23 @@ STATIC mp_obj_t array_subscr(mp_obj_t self_in, mp_obj_t index_in, mp_obj_t value
                     mp_get_buffer_raise(value, &bufinfo, MP_BUFFER_READ);
                     src_len = bufinfo.len;
                     src_items = bufinfo.buf;
+                } else if (MP_OBJ_IS_TYPE(value, &mp_type_memoryview)) {
+                    mp_obj_array_t *src_slice = value;
+                    if (item_sz != mp_binary_get_size('@', src_slice->typecode & TYPECODE_MASK, NULL)) {
+                        goto compat_error;
+                    }
+                    src_len = src_slice->len;
+                    src_items = (uint8_t*)src_slice->items + (src_slice->free * item_sz);
                 } else {
                     mp_not_implemented("array/bytes required on right side");
                 }
 
                 // TODO: check src/dst compat
                 mp_int_t len_adj = src_len - (slice.stop - slice.start);
                 if (len_adj > 0) {
+                    if (o->base.type == &mp_type_memoryview) {
+                        goto compat_error;
+                    }
                     if (len_adj > o->free) {
                         // TODO: alloc policy; at the moment we go conservative
                         o->items = m_renew(byte, o->items, (o->len + o->free) * item_sz, (o->len + len_adj) * item_sz);
@@ -405,12 +415,20 @@ STATIC mp_obj_t array_subscr(mp_obj_t self_in, mp_obj_t index_in, mp_obj_t value
                     mp_seq_replace_slice_grow_inplace(o->items, o->len,
                         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_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);
-                    // TODO: alloc policy after shrinking
+                    if (o->base.type == &mp_type_memoryview) {
+                        if (len_adj != 0) {
+                            goto compat_error;
+                        }
+                        mp_seq_replace_slice_no_grow((uint8_t*)o->items + (o->free * item_sz), o->len,
+                            slice.start, slice.stop, src_items, src_len, item_sz);
+                    } else {
+                        mp_seq_replace_slice_no_grow(o->items, o->len,
+                            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);
+                        // TODO: alloc policy after shrinking
+                    }
                 }
                 o->len += len_adj;
                 return mp_const_none;
diff --git a/tests/basics/memoryview1.py b/tests/basics/memoryview1.py
@@ -36,3 +36,45 @@
 print(list(m[1:-1]))
 m[2] = 6
 print(a)
+
+# test slice assignment between memoryviews
+b1 = bytearray(b'1234')
+b2 = bytearray(b'5678')
+b3 = bytearray(b'5678')
+m1 = memoryview(b1)
+m2 = memoryview(b2)
+m3 = memoryview(b3)
+m2[1:3] = m1[0:2]
+print(b2)
+b3[1:3] = m1[0:2]
+print(b3)
+m1[2:4] = b3[1:3]
+print(b1)
+
+try:
+    m2[1:3] = b1[0:4]
+except ValueError:
+    print("ValueError")
+
+try:
+    m2[1:3] = m1[0:4]
+except ValueError:
+    print("ValueError")
+
+try:
+    m2[0:4] = m1[1:3]
+except ValueError:
+    print("ValueError")
+
+# test memoryview of arrays with items sized larger than 1
+a1 = array.array('i', [0]*5)
+m4 = memoryview(a1)
+a2 = array.array('i', [3]*5)
+m5 = memoryview(a2)
+m4[1:3] = m5[1:3]
+print(a1)
+
+try:
+    m4[1:3] = m2[1:3]
+except ValueError:
+    print("ValueError")
