Merge pull request RustPython#3064 from qingshi163/array-pickle

Implement __reduce__ and __reduce_ex__ for array

Author: youknowone

Lib/test/test_array.py

@@ -247,17 +247,13 @@ def test_deepcopy(self):
         self.assertNotEqual(id(a), id(b))
         self.assertEqual(a, b)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_reduce_ex(self):
         a = array.array(self.typecode, self.example)
         for protocol in range(3):
             self.assertIs(a.__reduce_ex__(protocol)[0], array.array)
         for protocol in range(3, pickle.HIGHEST_PROTOCOL + 1):
             self.assertIs(a.__reduce_ex__(protocol)[0], array_reconstructor)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_pickle(self):
         for protocol in range(pickle.HIGHEST_PROTOCOL + 1):
             a = array.array(self.typecode, self.example)
@@ -273,8 +269,6 @@ def test_pickle(self):
             self.assertEqual(a.x, b.x)
             self.assertEqual(type(a), type(b))
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_pickle_for_empty_array(self):
         for protocol in range(pickle.HIGHEST_PROTOCOL + 1):
             a = array.array(self.typecode)

extra_tests/snippets/stdlib_array.py

@@ -1,5 +1,6 @@
 from testutils import assert_raises
 from array import array
+from pickle import dumps, loads
 
 a1 = array("b", [0, 1, 2, 3])
 
@@ -96,4 +97,11 @@ def test_array_frombytes():
 with assert_raises(IndexError):
     a[0] = 42
 with assert_raises(IndexError):
-    del a[42]
+    del a[42]
+
+test_str = '🌉abc🌐def🌉🌐'
+u = array('u', test_str)
+# skip as 2 bytes character enviroment with CPython is failing the test
+if u.itemsize >= 4:
+    assert u.__reduce_ex__(1)[1][1] == list(test_str)
+    assert loads(dumps(u, 1)) == loads(dumps(u, 3))

stdlib/src/array.rs

@@ -12,8 +12,8 @@ mod array {
     };
     use crate::vm::{
         builtins::{
-            PyByteArray, PyBytes, PyBytesRef, PyIntRef, PyList, PyListRef, PySliceRef, PyStr,
-            PyStrRef, PyTypeRef,
+            PyByteArray, PyBytes, PyBytesRef, PyDictRef, PyFloat, PyInt, PyIntRef, PyList,
+            PyListRef, PySliceRef, PyStr, PyStrRef, PyTypeRef,
         },
         class_or_notimplemented,
         function::{
@@ -28,8 +28,8 @@ mod array {
             AsBuffer, AsMapping, Comparable, Iterable, IteratorIterable, PyComparisonOp,
             SlotConstructor, SlotIterator,
         },
-        IdProtocol, PyComparisonValue, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
-        TypeProtocol, VirtualMachine,
+        IdProtocol, PyComparisonValue, PyObjectRef, PyObjectWrap, PyRef, PyResult, PyValue,
+        TryFromObject, TypeProtocol, VirtualMachine,
     };
     use crossbeam_utils::atomic::AtomicCell;
     use itertools::Itertools;
@@ -460,6 +460,14 @@ mod array {
                         })*
                     }
                 }
+
+                fn get_objects(&self, vm: &VirtualMachine) -> Vec<PyObjectRef> {
+                    match self {
+                        $(ArrayContentType::$n(v) => {
+                            v.iter().map(|&x| x.to_object(vm)).collect()
+                        })*
+                    }
+                }
             }
         };
     }
@@ -486,32 +494,41 @@ mod array {
     trait ArrayElement: Sized {
         fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self>;
         fn byteswap(self) -> Self;
+        fn to_object(self, vm: &VirtualMachine) -> PyObjectRef;
     }
 
     macro_rules! impl_array_element {
-        ($(($t:ty, $f_into:path, $f_swap:path),)*) => {$(
+        ($(($t:ty, $f_from:path, $f_swap:path, $f_to:path),)*) => {$(
             impl ArrayElement for $t {
                 fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
-                    $f_into(vm, obj)
+                    $f_from(vm, obj)
                 }
                 fn byteswap(self) -> Self {
                     $f_swap(self)
                 }
+                fn to_object(self, vm: &VirtualMachine) -> PyObjectRef {
+                    $f_to(self).into_object(vm)
+                }
             }
         )*};
     }
 
     impl_array_element!(
-        (i8, i8::try_from_object, i8::swap_bytes),
-        (u8, u8::try_from_object, u8::swap_bytes),
-        (i16, i16::try_from_object, i16::swap_bytes),
-        (u16, u16::try_from_object, u16::swap_bytes),
-        (i32, i32::try_from_object, i32::swap_bytes),
-        (u32, u32::try_from_object, u32::swap_bytes),
-        (i64, i64::try_from_object, i64::swap_bytes),
-        (u64, u64::try_from_object, u64::swap_bytes),
-        (f32, f32_try_into_from_object, f32_swap_bytes),
-        (f64, f64_try_into_from_object, f64_swap_bytes),
+        (i8, i8::try_from_object, i8::swap_bytes, PyInt::from),
+        (u8, u8::try_from_object, u8::swap_bytes, PyInt::from),
+        (i16, i16::try_from_object, i16::swap_bytes, PyInt::from),
+        (u16, u16::try_from_object, u16::swap_bytes, PyInt::from),
+        (i32, i32::try_from_object, i32::swap_bytes, PyInt::from),
+        (u32, u32::try_from_object, u32::swap_bytes, PyInt::from),
+        (i64, i64::try_from_object, i64::swap_bytes, PyInt::from),
+        (u64, u64::try_from_object, u64::swap_bytes, PyInt::from),
+        (
+            f32,
+            f32_try_into_from_object,
+            f32_swap_bytes,
+            pyfloat_from_f32
+        ),
+        (f64, f64_try_into_from_object, f64_swap_bytes, PyFloat::from),
     );
 
     fn f32_swap_bytes(x: f32) -> f32 {
@@ -530,6 +547,10 @@ mod array {
         ArgIntoFloat::try_from_object(vm, obj).map(|x| x.to_f64())
     }
 
+    fn pyfloat_from_f32(value: f32) -> PyFloat {
+        PyFloat::from(value as f64)
+    }
+
     impl ArrayElement for WideChar {
         fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
             PyStrRef::try_from_object(vm, obj)?
@@ -542,6 +563,9 @@ mod array {
         fn byteswap(self) -> Self {
             Self(self.0.swap_bytes())
         }
+        fn to_object(self, _vm: &VirtualMachine) -> PyObjectRef {
+            unreachable!()
+        }
     }
 
     fn u32_to_char(ch: u32) -> Result<char, String> {
@@ -731,6 +755,38 @@ mod array {
             }
         }
 
+        fn _wchar_bytes_to_string(
+            bytes: &[u8],
+            item_size: usize,
+            vm: &VirtualMachine,
+        ) -> PyResult<String> {
+            if item_size == 2 {
+                // safe because every configuration of bytes for the types we support are valid
+                let utf16 = unsafe {
+                    std::slice::from_raw_parts(
+                        bytes.as_ptr() as *const u16,
+                        bytes.len() / std::mem::size_of::<u16>(),
+                    )
+                };
+                Ok(String::from_utf16_lossy(utf16))
+            } else {
+                // safe because every configuration of bytes for the types we support are valid
+                let chars = unsafe {
+                    std::slice::from_raw_parts(
+                        bytes.as_ptr() as *const u32,
+                        bytes.len() / std::mem::size_of::<u32>(),
+                    )
+                };
+                chars
+                    .iter()
+                    .map(|&ch| {
+                        // cpython issue 17223
+                        u32_to_char(ch).map_err(|msg| vm.new_value_error(msg))
+                    })
+                    .try_collect()
+            }
+        }
+
         fn _unicode_to_wchar_bytes(utf8: &str, item_size: usize) -> Vec<u8> {
             if item_size == 2 {
                 utf8.encode_utf16()
@@ -771,31 +827,7 @@ mod array {
                 ));
             }
             let bytes = array.get_bytes();
-            if self.itemsize() == 2 {
-                // safe because every configuration of bytes for the types we support are valid
-                let utf16 = unsafe {
-                    std::slice::from_raw_parts(
-                        bytes.as_ptr() as *const u16,
-                        bytes.len() / std::mem::size_of::<u16>(),
-                    )
-                };
-                Ok(String::from_utf16_lossy(utf16))
-            } else {
-                // safe because every configuration of bytes for the types we support are valid
-                let chars = unsafe {
-                    std::slice::from_raw_parts(
-                        bytes.as_ptr() as *const u32,
-                        bytes.len() / std::mem::size_of::<u32>(),
-                    )
-                };
-                chars
-                    .iter()
-                    .map(|&ch| {
-                        // cpython issue 17223
-                        u32_to_char(ch).map_err(|msg| vm.new_value_error(msg))
-                    })
-                    .try_collect()
-            }
+            Self::_wchar_bytes_to_string(bytes, self.itemsize(), vm)
         }
 
         fn _from_bytes(&self, b: &[u8], itemsize: usize, vm: &VirtualMachine) -> PyResult<()> {
@@ -1079,6 +1111,52 @@ mod array {
             }
             Ok(true)
         }
+
+        #[pymethod(magic)]
+        fn reduce_ex(
+            zelf: PyRef<Self>,
+            proto: usize,
+            vm: &VirtualMachine,
+        ) -> PyResult<(PyObjectRef, PyObjectRef, Option<PyDictRef>)> {
+            if proto < 3 {
+                return Self::reduce(zelf, vm);
+            }
+            let array = zelf.read();
+            let cls = zelf.as_object().clone_class().into_object();
+            let typecode = vm.ctx.new_utf8_str(array.typecode_str());
+            let bytes = vm.ctx.new_bytes(array.get_bytes().to_vec());
+            let code = MachineFormatCode::from_typecode(array.typecode()).unwrap();
+            let code = PyInt::from(u8::from(code)).into_object(vm);
+            let module = vm.import("array", None, 0)?;
+            let func = vm.get_attribute(module, "_array_reconstructor")?;
+            Ok((
+                func,
+                vm.ctx.new_tuple(vec![cls, typecode, code, bytes]),
+                zelf.as_object().dict(),
+            ))
+        }
+
+        #[pymethod(magic)]
+        fn reduce(
+            zelf: PyRef<Self>,
+            vm: &VirtualMachine,
+        ) -> PyResult<(PyObjectRef, PyObjectRef, Option<PyDictRef>)> {
+            let array = zelf.read();
+            let cls = zelf.as_object().clone_class().into_object();
+            let typecode = vm.ctx.new_utf8_str(array.typecode_str());
+            let values = if array.typecode() == 'u' {
+                let s = Self::_wchar_bytes_to_string(array.get_bytes(), array.itemsize(), vm)?;
+                s.chars().map(|x| x.into_pyobject(vm)).collect()
+            } else {
+                array.get_objects(vm)
+            };
+            let values = vm.ctx.new_list(values);
+            Ok((
+                cls,
+                vm.ctx.new_tuple(vec![typecode, values]),
+                zelf.as_object().dict(),
+            ))
+        }
     }
 
     impl Comparable for PyArray {