Skip to content

Commit 7595e67

Browse files
authored
gh-120380: fix Python implementation of pickle.Pickler for bytes and bytearray objects in protocol version 5. (GH-120422)
1 parent 83d3d7a commit 7595e67

File tree

3 files changed

+81
-21
lines changed

3 files changed

+81
-21
lines changed

Lib/pickle.py

+35-15
Original file line numberDiff line numberDiff line change
@@ -782,14 +782,10 @@ def save_float(self, obj):
782782
self.write(FLOAT + repr(obj).encode("ascii") + b'\n')
783783
dispatch[float] = save_float
784784

785-
def save_bytes(self, obj):
786-
if self.proto < 3:
787-
if not obj: # bytes object is empty
788-
self.save_reduce(bytes, (), obj=obj)
789-
else:
790-
self.save_reduce(codecs.encode,
791-
(str(obj, 'latin1'), 'latin1'), obj=obj)
792-
return
785+
def _save_bytes_no_memo(self, obj):
786+
# helper for writing bytes objects for protocol >= 3
787+
# without memoizing them
788+
assert self.proto >= 3
793789
n = len(obj)
794790
if n <= 0xff:
795791
self.write(SHORT_BINBYTES + pack("<B", n) + obj)
@@ -799,21 +795,37 @@ def save_bytes(self, obj):
799795
self._write_large_bytes(BINBYTES + pack("<I", n), obj)
800796
else:
801797
self.write(BINBYTES + pack("<I", n) + obj)
798+
799+
def save_bytes(self, obj):
800+
if self.proto < 3:
801+
if not obj: # bytes object is empty
802+
self.save_reduce(bytes, (), obj=obj)
803+
else:
804+
self.save_reduce(codecs.encode,
805+
(str(obj, 'latin1'), 'latin1'), obj=obj)
806+
return
807+
self._save_bytes_no_memo(obj)
802808
self.memoize(obj)
803809
dispatch[bytes] = save_bytes
804810

811+
def _save_bytearray_no_memo(self, obj):
812+
# helper for writing bytearray objects for protocol >= 5
813+
# without memoizing them
814+
assert self.proto >= 5
815+
n = len(obj)
816+
if n >= self.framer._FRAME_SIZE_TARGET:
817+
self._write_large_bytes(BYTEARRAY8 + pack("<Q", n), obj)
818+
else:
819+
self.write(BYTEARRAY8 + pack("<Q", n) + obj)
820+
805821
def save_bytearray(self, obj):
806822
if self.proto < 5:
807823
if not obj: # bytearray is empty
808824
self.save_reduce(bytearray, (), obj=obj)
809825
else:
810826
self.save_reduce(bytearray, (bytes(obj),), obj=obj)
811827
return
812-
n = len(obj)
813-
if n >= self.framer._FRAME_SIZE_TARGET:
814-
self._write_large_bytes(BYTEARRAY8 + pack("<Q", n), obj)
815-
else:
816-
self.write(BYTEARRAY8 + pack("<Q", n) + obj)
828+
self._save_bytearray_no_memo(obj)
817829
self.memoize(obj)
818830
dispatch[bytearray] = save_bytearray
819831

@@ -832,10 +844,18 @@ def save_picklebuffer(self, obj):
832844
if in_band:
833845
# Write data in-band
834846
# XXX The C implementation avoids a copy here
847+
buf = m.tobytes()
848+
in_memo = id(buf) in self.memo
835849
if m.readonly:
836-
self.save_bytes(m.tobytes())
850+
if in_memo:
851+
self._save_bytes_no_memo(buf)
852+
else:
853+
self.save_bytes(buf)
837854
else:
838-
self.save_bytearray(m.tobytes())
855+
if in_memo:
856+
self._save_bytearray_no_memo(buf)
857+
else:
858+
self.save_bytearray(buf)
839859
else:
840860
# Write data out-of-band
841861
self.write(NEXT_BUFFER)

Lib/test/pickletester.py

+43-6
Original file line numberDiff line numberDiff line change
@@ -1845,6 +1845,25 @@ def test_bytes(self):
18451845
p = self.dumps(s, proto)
18461846
self.assert_is_copy(s, self.loads(p))
18471847

1848+
def test_bytes_memoization(self):
1849+
for proto in protocols:
1850+
for array_type in [bytes, ZeroCopyBytes]:
1851+
for s in b'', b'xyz', b'xyz'*100:
1852+
with self.subTest(proto=proto, array_type=array_type, s=s, independent=False):
1853+
b = array_type(s)
1854+
p = self.dumps((b, b), proto)
1855+
x, y = self.loads(p)
1856+
self.assertIs(x, y)
1857+
self.assert_is_copy((b, b), (x, y))
1858+
1859+
with self.subTest(proto=proto, array_type=array_type, s=s, independent=True):
1860+
b1, b2 = array_type(s), array_type(s)
1861+
p = self.dumps((b1, b2), proto)
1862+
# Note that (b1, b2) = self.loads(p) might have identical
1863+
# components, i.e., b1 is b2, but this is not always the
1864+
# case if the content is large (equality still holds).
1865+
self.assert_is_copy((b1, b2), self.loads(p))
1866+
18481867
def test_bytearray(self):
18491868
for proto in protocols:
18501869
for s in b'', b'xyz', b'xyz'*100:
@@ -1864,13 +1883,31 @@ def test_bytearray(self):
18641883
self.assertNotIn(b'bytearray', p)
18651884
self.assertTrue(opcode_in_pickle(pickle.BYTEARRAY8, p))
18661885

1867-
def test_bytearray_memoization_bug(self):
1886+
def test_bytearray_memoization(self):
18681887
for proto in protocols:
1869-
for s in b'', b'xyz', b'xyz'*100:
1870-
b = bytearray(s)
1871-
p = self.dumps((b, b), proto)
1872-
b1, b2 = self.loads(p)
1873-
self.assertIs(b1, b2)
1888+
for array_type in [bytearray, ZeroCopyBytearray]:
1889+
for s in b'', b'xyz', b'xyz'*100:
1890+
with self.subTest(proto=proto, array_type=array_type, s=s, independent=False):
1891+
b = array_type(s)
1892+
p = self.dumps((b, b), proto)
1893+
b1, b2 = self.loads(p)
1894+
self.assertIs(b1, b2)
1895+
1896+
with self.subTest(proto=proto, array_type=array_type, s=s, independent=True):
1897+
b1a, b2a = array_type(s), array_type(s)
1898+
# Unlike bytes, equal but independent bytearray objects are
1899+
# never identical.
1900+
self.assertIsNot(b1a, b2a)
1901+
1902+
p = self.dumps((b1a, b2a), proto)
1903+
b1b, b2b = self.loads(p)
1904+
self.assertIsNot(b1b, b2b)
1905+
1906+
self.assertIsNot(b1a, b1b)
1907+
self.assert_is_copy(b1a, b1b)
1908+
1909+
self.assertIsNot(b2a, b2b)
1910+
self.assert_is_copy(b2a, b2b)
18741911

18751912
def test_ints(self):
18761913
for proto in protocols:
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Fix Python implementation of :class:`pickle.Pickler` for :class:`bytes` and
2+
:class:`bytearray` objects when using protocol version 5. Patch by Bénédikt
3+
Tran.

0 commit comments

Comments
 (0)