Merge pull request RustPython#2383 from qingshi163/memoryview-cast

Implement Memoryview cast
AlanWZHU · Jan 3, 2021 · b0183b2 · b0183b2
2 parents 0786e4c + 06ff997
commit b0183b2
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 2 deletions.
diff --git a/Lib/test/test_memoryview.py b/Lib/test/test_memoryview.py
@@ -384,8 +384,6 @@ def test_toreadonly(self):
             mm.release()
             m.tolist()
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_issue22668(self):
         a = array.array('H', [256, 256, 256, 256])
         x = memoryview(a)

diff --git a/vm/src/builtins/memory.rs b/vm/src/builtins/memory.rs
@@ -23,6 +23,7 @@ use crossbeam_utils::atomic::AtomicCell;
 use itertools::Itertools;
 use num_bigint::BigInt;
 use num_traits::{One, Signed, ToPrimitive, Zero};
+use rustpython_common::borrow::BorrowValue;
 
 #[derive(Debug)]
 pub struct BufferRef(Box<dyn Buffer>);
@@ -654,6 +655,7 @@ impl PyMemoryView {
     fn toreadonly(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
         zelf.try_not_released(vm)?;
         let buffer = BufferRef(Box::new(zelf.clone()));
+        zelf.exports.fetch_add(1);
         Ok(PyMemoryView {
             obj: zelf.obj.clone(),
             buffer,
@@ -703,6 +705,48 @@ impl PyMemoryView {
         bytes_to_hex(bytes, sep, bytes_per_sep, vm)
     }
 
+    // TODO: support cast shape
+    #[pymethod]
+    fn cast(zelf: PyRef<Self>, format: PyStrRef, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
+        zelf.try_not_released(vm)?;
+        if !zelf.options.contiguous {
+            return Err(vm.new_type_error(
+                "memoryview: casts are restricted to C-contiguous views".to_owned(),
+            ));
+        }
+
+        let format_spec = Self::parse_format(format.borrow_value(), vm)?;
+        let itemsize = format_spec.size();
+        let bytelen = zelf.options.len * zelf.options.itemsize;
+
+        if bytelen % itemsize != 0 {
+            return Err(
+                vm.new_type_error("memoryview: length is not a multiple of itemsize".to_owned())
+            );
+        }
+
+        let buffer = BufferRef(Box::new(zelf.clone()));
+        zelf.exports.fetch_add(1);
+
+        Ok(PyMemoryView {
+            obj: zelf.obj.clone(),
+            buffer,
+            options: BufferOptions {
+                itemsize,
+                len: bytelen / itemsize,
+                format: format.to_string().into(),
+                ..zelf.options.clone()
+            },
+            released: AtomicCell::new(false),
+            stop: zelf.stop + itemsize - zelf.options.itemsize,
+            exports: AtomicCell::new(0),
+            format_spec,
+            hash: OnceCell::new(),
+            ..*zelf
+        }
+        .into_ref(vm))
+    }
+
     fn eq(zelf: &PyRef<Self>, other: &PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
         if zelf.is(other) {
             return Ok(true);