Merge pull request RustPython#2302 from qingshi163/bytes_pickle

Fix bytes pickle with subclass

Author: coolreader18

Lib/test/test_bytes.py

@@ -1860,8 +1860,6 @@ def test_join(self):
         s3 = s1.join([b"abcd"])
         self.assertIs(type(s3), self.basetype)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_pickle(self):
         a = self.type2test(b"abcd")
         a.x = 10
@@ -1875,8 +1873,6 @@ def test_pickle(self):
             self.assertEqual(type(a), type(b))
             self.assertEqual(type(a.y), type(b.y))
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_copy(self):
         a = self.type2test(b"abcd")
         a.x = 10

extra_tests/snippets/bytearray.py

@@ -729,3 +729,28 @@
 assert pickle.dumps(a, 2) == b'\x80\x02c__builtin__\nbytearray\nq\x00c_codecs\nencode\nq\x01X\x0c\x00\x00\x00\xc3\xbfab\xc2\x80\x00\x00\xc3\xb8\x00\x00q\x02X\x06\x00\x00\x00latin1q\x03\x86q\x04Rq\x05\x85q\x06Rq\x07.'
 assert pickle.dumps(a, 3) == b'\x80\x03cbuiltins\nbytearray\nq\x00C\t\xffab\x80\x00\x00\xf8\x00\x00q\x01\x85q\x02Rq\x03.'
 assert pickle.dumps(a, 4) == b'\x80\x04\x95*\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x94\x8c\tbytearray\x94\x93\x94C\t\xffab\x80\x00\x00\xf8\x00\x00\x94\x85\x94R\x94.'
+
+# pickle with subclass
+class A(bytes):
+    pass
+
+a = A()
+a.x = 10
+a.y = A(b'123')
+b = pickle.loads(pickle.dumps(a, 4))
+assert type(a) == type(b)
+assert a.x == b.x
+assert a.y == b.y
+assert a == b
+
+class B(bytearray):
+    pass
+
+a = B()
+a.x = 10
+a.y = B(b'123')
+b = pickle.loads(pickle.dumps(a, 4))
+assert type(a) == type(b)
+assert a.x == b.x
+assert a.y == b.y
+assert a == b

vm/src/builtins/bytearray.rs

@@ -1,5 +1,6 @@
 //! Implementation of the python bytearray object.
 use super::bytes::{PyBytes, PyBytesRef};
+use super::dict::PyDictRef;
 use super::int::PyIntRef;
 use super::memory::{Buffer, BufferOptions, ResizeGuard};
 use super::pystr::PyStrRef;
@@ -19,7 +20,7 @@ use crate::common::lock::{
 use crate::function::{OptionalArg, OptionalOption};
 use crate::pyobject::{
     BorrowValue, Either, IdProtocol, IntoPyObject, PyClassImpl, PyComparisonValue, PyContext,
-    PyIterable, PyObjectRef, PyRef, PyResult, PyValue,
+    PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
 };
 use crate::sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex};
 use crate::slots::{BufferProtocol, Comparable, Hashable, PyComparisonOp, Unhashable};
@@ -627,16 +628,24 @@ impl PyByteArray {
     }
 
     #[pymethod(magic)]
-    fn reduce_ex(zelf: PyRef<Self>, _proto: usize, vm: &VirtualMachine) -> (PyTypeRef, PyTupleRef) {
+    fn reduce_ex(
+        zelf: PyRef<Self>,
+        _proto: usize,
+        vm: &VirtualMachine,
+    ) -> (PyTypeRef, PyTupleRef, Option<PyDictRef>) {
         Self::reduce(zelf, vm)
     }
 
     #[pymethod(magic)]
-    fn reduce(zelf: PyRef<Self>, vm: &VirtualMachine) -> (PyTypeRef, PyTupleRef) {
+    fn reduce(
+        zelf: PyRef<Self>,
+        vm: &VirtualMachine,
+    ) -> (PyTypeRef, PyTupleRef, Option<PyDictRef>) {
         let bytes = PyBytes::from(zelf.borrow_value().elements.clone()).into_pyobject(vm);
         (
-            Self::class(vm).clone(),
+            zelf.as_object().clone_class(),
             PyTupleRef::with_elements(vec![bytes], &vm.ctx),
+            zelf.as_object().dict(),
         )
     }
 }

vm/src/builtins/bytes.rs

@@ -7,6 +7,7 @@ use rustpython_common::{
 use std::mem::size_of;
 use std::ops::Deref;
 
+use super::dict::PyDictRef;
 use super::int::PyIntRef;
 use super::pystr::PyStrRef;
 use super::pytype::PyTypeRef;
@@ -19,7 +20,7 @@ use crate::common::hash::PyHash;
 use crate::function::{OptionalArg, OptionalOption};
 use crate::pyobject::{
     BorrowValue, Either, IntoPyObject, PyClassImpl, PyComparisonValue, PyContext, PyIterable,
-    PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
+    PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol,
 };
 use crate::slots::{BufferProtocol, Comparable, Hashable, PyComparisonOp};
 use crate::vm::VirtualMachine;
@@ -472,6 +473,28 @@ impl PyBytes {
             .collect();
         PyTupleRef::with_elements(param, &vm.ctx)
     }
+
+    #[pymethod(magic)]
+    fn reduce_ex(
+        zelf: PyRef<Self>,
+        _proto: usize,
+        vm: &VirtualMachine,
+    ) -> (PyTypeRef, PyTupleRef, Option<PyDictRef>) {
+        Self::reduce(zelf, vm)
+    }
+
+    #[pymethod(magic)]
+    fn reduce(
+        zelf: PyRef<Self>,
+        vm: &VirtualMachine,
+    ) -> (PyTypeRef, PyTupleRef, Option<PyDictRef>) {
+        let bytes = PyBytes::from(zelf.inner.elements.clone()).into_pyobject(vm);
+        (
+            zelf.as_object().clone_class(),
+            PyTupleRef::with_elements(vec![bytes], &vm.ctx),
+            zelf.as_object().dict(),
+        )
+    }
 }
 
 impl BufferProtocol for PyBytes {