From 36dd7396fcd26d8bf9919d536d05d7000becbe5b Mon Sep 17 00:00:00 2001 From: Ma Lin Date: Tue, 8 Mar 2022 17:33:56 +0800 Subject: bpo-44439: _ZipWriteFile.write() handle buffer protocol correctly (GH-29468) Co-authored-by: Marco Ribeiro --- Lib/test/test_zipfile.py | 9 +++++++++ Lib/zipfile.py | 9 ++++++++- .../NEWS.d/next/Library/2021-11-08-20-27-41.bpo-44439.I_8qro.rst | 2 ++ 3 files changed, 19 insertions(+), 1 deletion(-) create mode 100644 Misc/NEWS.d/next/Library/2021-11-08-20-27-41.bpo-44439.I_8qro.rst diff --git a/Lib/test/test_zipfile.py b/Lib/test/test_zipfile.py index de2dd33..759a4ab 100644 --- a/Lib/test/test_zipfile.py +++ b/Lib/test/test_zipfile.py @@ -1,3 +1,4 @@ +import array import contextlib import importlib.util import io @@ -1121,6 +1122,14 @@ class AbstractWriterTests: self.assertRaises(ValueError, w.write, b'') self.assertEqual(zipf.read('test'), data) + def test_issue44439(self): + q = array.array('Q', [1, 2, 3, 4, 5]) + LENGTH = len(q) * q.itemsize + with zipfile.ZipFile(io.BytesIO(), 'w', self.compression) as zip: + with zip.open('data', 'w') as data: + self.assertEqual(data.write(q), LENGTH) + self.assertEqual(zip.getinfo('data').file_size, LENGTH) + class StoredWriterTests(AbstractWriterTests, unittest.TestCase): compression = zipfile.ZIP_STORED diff --git a/Lib/zipfile.py b/Lib/zipfile.py index 8e9325b..41bf49a 100644 --- a/Lib/zipfile.py +++ b/Lib/zipfile.py @@ -1147,8 +1147,15 @@ class _ZipWriteFile(io.BufferedIOBase): def write(self, data): if self.closed: raise ValueError('I/O operation on closed file.') - nbytes = len(data) + + # Accept any data that supports the buffer protocol + if isinstance(data, (bytes, bytearray)): + nbytes = len(data) + else: + data = memoryview(data) + nbytes = data.nbytes self._file_size += nbytes + self._crc = crc32(data, self._crc) if self._compressor: data = self._compressor.compress(data) diff --git a/Misc/NEWS.d/next/Library/2021-11-08-20-27-41.bpo-44439.I_8qro.rst b/Misc/NEWS.d/next/Library/2021-11-08-20-27-41.bpo-44439.I_8qro.rst new file mode 100644 index 0000000..f4e562c --- /dev/null +++ b/Misc/NEWS.d/next/Library/2021-11-08-20-27-41.bpo-44439.I_8qro.rst @@ -0,0 +1,2 @@ +Fix ``.write()`` method of a member file in ``ZipFile``, when the input data is +an object that supports the buffer protocol, the file length may be wrong. -- cgit v0.12