[Python-checkins] bpo-44439: BZ2File.write() / LZMAFile.write() handle buffer protocol correctly (GH-26764)

serhiy-storchaka webhook-mailer at python.org
Tue Jun 22 03:04:48 EDT 2021


https://github.com/python/cpython/commit/bc6c12c72a9536acc96e7b9355fd69d1083a43c1
commit: bc6c12c72a9536acc96e7b9355fd69d1083a43c1
branch: main
author: Ma Lin <animalize at users.noreply.github.com>
committer: serhiy-storchaka <storchaka at gmail.com>
date: 2021-06-22T10:04:23+03:00
summary:

bpo-44439: BZ2File.write() / LZMAFile.write() handle buffer protocol correctly (GH-26764)

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.

files:
A Misc/NEWS.d/next/Library/2021-06-17-15-01-51.bpo-44439.1S7QhT.rst
M Lib/bz2.py
M Lib/gzip.py
M Lib/lzma.py
M Lib/test/test_bz2.py
M Lib/test/test_gzip.py
M Lib/test/test_lzma.py

diff --git a/Lib/bz2.py b/Lib/bz2.py
index a2c588e7487f3..7f1d20632ef13 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.
diff --git a/Lib/gzip.py b/Lib/gzip.py
index 1c1e795e1715d..3d837b744800e 100644
--- a/Lib/gzip.py
+++ b/Lib/gzip.py
@@ -278,7 +278,7 @@ def write(self,data):
         if self.fileobj is None:
             raise ValueError("write() on closed GzipFile object")
 
-        if isinstance(data, bytes):
+        if isinstance(data, (bytes, bytearray)):
             length = len(data)
         else:
             # accept any data that supports the buffer protocol
diff --git a/Lib/lzma.py b/Lib/lzma.py
index 2ada7d81d3c81..9abf06d91db18 100644
--- a/Lib/lzma.py
+++ b/Lib/lzma.py
@@ -229,14 +229,22 @@ def write(self, data):
         """Write a bytes object 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 seek(self, offset, whence=io.SEEK_SET):
         """Change the file position.
diff --git a/Lib/test/test_bz2.py b/Lib/test/test_bz2.py
index efed3a859ba21..7913beb87a352 100644
--- a/Lib/test/test_bz2.py
+++ b/Lib/test/test_bz2.py
@@ -1,6 +1,7 @@
 from test import support
 from test.support import bigmemtest, _4G
 
+import array
 import unittest
 from io import BytesIO, DEFAULT_BUFFER_SIZE
 import os
@@ -620,6 +621,14 @@ def test_read_truncated(self):
             with BZ2File(BytesIO(truncated[:i])) as f:
                 self.assertRaises(EOFError, f.read, 1)
 
+    def test_issue44439(self):
+        q = array.array('Q', [1, 2, 3, 4, 5])
+        LENGTH = len(q) * q.itemsize
+
+        with BZ2File(BytesIO(), 'w') as f:
+            self.assertEqual(f.write(q), LENGTH)
+            self.assertEqual(f.tell(), LENGTH)
+
 
 class BZ2CompressorTest(BaseTest):
     def testCompress(self):
diff --git a/Lib/test/test_gzip.py b/Lib/test/test_gzip.py
index 446b61ab439ff..7b51e45aad92b 100644
--- a/Lib/test/test_gzip.py
+++ b/Lib/test/test_gzip.py
@@ -592,6 +592,15 @@ def test_prepend_error(self):
         with gzip.open(self.filename, "rb") as f:
             f._buffer.raw._fp.prepend()
 
+    def test_issue44439(self):
+        q = array.array('Q', [1, 2, 3, 4, 5])
+        LENGTH = len(q) * q.itemsize
+
+        with gzip.GzipFile(fileobj=io.BytesIO(), mode='w') as f:
+            self.assertEqual(f.write(q), LENGTH)
+            self.assertEqual(f.tell(), LENGTH)
+
+
 class TestOpen(BaseTest):
     def test_binary_modes(self):
         uncompressed = data1 * 50
diff --git a/Lib/test/test_lzma.py b/Lib/test/test_lzma.py
index db20300056e48..1e2066b89168f 100644
--- a/Lib/test/test_lzma.py
+++ b/Lib/test/test_lzma.py
@@ -1,4 +1,5 @@
 import _compression
+import array
 from io import BytesIO, UnsupportedOperation, DEFAULT_BUFFER_SIZE
 import os
 import pathlib
@@ -1231,6 +1232,14 @@ def test_issue21872(self):
         self.assertTrue(d2.eof)
         self.assertEqual(out1 + out2, entire)
 
+    def test_issue44439(self):
+        q = array.array('Q', [1, 2, 3, 4, 5])
+        LENGTH = len(q) * q.itemsize
+
+        with LZMAFile(BytesIO(), 'w') as f:
+            self.assertEqual(f.write(q), LENGTH)
+            self.assertEqual(f.tell(), LENGTH)
+
 
 class OpenTestCase(unittest.TestCase):
 
diff --git a/Misc/NEWS.d/next/Library/2021-06-17-15-01-51.bpo-44439.1S7QhT.rst b/Misc/NEWS.d/next/Library/2021-06-17-15-01-51.bpo-44439.1S7QhT.rst
new file mode 100644
index 0000000000000..27396683700a8
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2021-06-17-15-01-51.bpo-44439.1S7QhT.rst
@@ -0,0 +1,3 @@
+Fix in :meth:`bz2.BZ2File.write` / :meth:`lzma.LZMAFile.write` methods, when
+the input data is an object that supports the buffer protocol, the file length
+may be wrong.



More information about the Python-checkins mailing list