[Python-checkins] bpo-44439: _ZipWriteFile.write() handle buffer protocol correctly (GH-29468)

[serhiy-storchaka]

https://github.com/python/cpython/commit/36dd7396fcd26d8bf9919d536d05d7000becbe5b
commit: 36dd7396fcd26d8bf9919d536d05d7000becbe5b
branch: main
author: Ma Lin <animalize at users.noreply.github.com>
committer: serhiy-storchaka <storchaka at gmail.com>
date: 2022-03-08T11:33:56+02:00
summary:

bpo-44439: _ZipWriteFile.write() handle buffer protocol correctly (GH-29468)

Co-authored-by: Marco Ribeiro <marcoffee at users.noreply.github.com>

files:
A Misc/NEWS.d/next/Library/2021-11-08-20-27-41.bpo-44439.I_8qro.rst
M Lib/test/test_zipfile.py
M Lib/zipfile.py

diff --git a/Lib/test/test_zipfile.py b/Lib/test/test_zipfile.py
index de2dd33f43660..759a4abb9d4d4 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 @@ def test_write_after_close(self):
             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 8e9325b934326..41bf49a8fe685 100644
--- a/Lib/zipfile.py
+++ b/Lib/zipfile.py
@@ -1147,8 +1147,15 @@ def writable(self):
     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 0000000000000..f4e562c4236d2
--- /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.