bpo-44439: BZ2File.write() / LZMAFile.write() handle buffer protocol correctly (GH-26764) (GH-26845)
No longer use len() to get the length of the input data. For some buffer protocol objects,
the length obtained by using len() is wrong.
(cherry picked from commit bc6c12c72a9536acc96e7b9355fd69d1083a43c1)
Co-authored-by: Ma Lin <animalize@users.noreply.github.com>
diff --git a/Lib/bz2.py b/Lib/bz2.py
index a2c588e..7f1d206 100644
--- a/Lib/bz2.py
+++ b/Lib/bz2.py
@@ -219,14 +219,22 @@ def write(self, data):
"""Write a byte string to the file.
Returns the number of uncompressed bytes written, which is
- always len(data). Note that due to buffering, the file on disk
- may not reflect the data written until close() is called.
+ always the length of data in bytes. Note that due to buffering,
+ the file on disk may not reflect the data written until close()
+ is called.
"""
self._check_can_write()
+ if isinstance(data, (bytes, bytearray)):
+ length = len(data)
+ else:
+ # accept any data that supports the buffer protocol
+ data = memoryview(data)
+ length = data.nbytes
+
compressed = self._compressor.compress(data)
self._fp.write(compressed)
- self._pos += len(data)
- return len(data)
+ self._pos += length
+ return length
def writelines(self, seq):
"""Write a sequence of byte strings to the file.