diff --git a/Lib/test/test_zipfile.py b/Lib/test/test_zipfile.py index 7b8922f4e05..14e1e08c5bf 100644 --- a/Lib/test/test_zipfile.py +++ b/Lib/test/test_zipfile.py @@ -402,6 +402,43 @@ class AbstractTestsWithSourceFile: self.assertEqual(one_info._compresslevel, 1) self.assertEqual(nine_info._compresslevel, 9) + def test_writing_errors(self): + class BrokenFile(io.BytesIO): + def write(self, data): + nonlocal count + if count is not None: + if count == stop: + raise OSError + count += 1 + super().write(data) + + stop = 0 + while True: + testfile = BrokenFile() + count = None + with zipfile.ZipFile(testfile, 'w', self.compression) as zipfp: + with zipfp.open('file1', 'w') as f: + f.write(b'data1') + count = 0 + try: + with zipfp.open('file2', 'w') as f: + f.write(b'data2') + except OSError: + stop += 1 + else: + break + finally: + count = None + with zipfile.ZipFile(io.BytesIO(testfile.getvalue())) as zipfp: + self.assertEqual(zipfp.namelist(), ['file1']) + self.assertEqual(zipfp.read('file1'), b'data1') + + with zipfile.ZipFile(io.BytesIO(testfile.getvalue())) as zipfp: + self.assertEqual(zipfp.namelist(), ['file1', 'file2']) + self.assertEqual(zipfp.read('file1'), b'data1') + self.assertEqual(zipfp.read('file2'), b'data2') + + def tearDown(self): unlink(TESTFN) unlink(TESTFN2) diff --git a/Lib/zipfile.py b/Lib/zipfile.py index 61cd929f614..2dc01647211 100644 --- a/Lib/zipfile.py +++ b/Lib/zipfile.py @@ -1105,47 +1105,50 @@ class _ZipWriteFile(io.BufferedIOBase): def close(self): if self.closed: return - super().close() - # Flush any data from the compressor, and update header info - if self._compressor: - buf = self._compressor.flush() - self._compress_size += len(buf) - self._fileobj.write(buf) - self._zinfo.compress_size = self._compress_size - else: - self._zinfo.compress_size = self._file_size - self._zinfo.CRC = self._crc - self._zinfo.file_size = self._file_size + try: + super().close() + # Flush any data from the compressor, and update header info + if self._compressor: + buf = self._compressor.flush() + self._compress_size += len(buf) + self._fileobj.write(buf) + self._zinfo.compress_size = self._compress_size + else: + self._zinfo.compress_size = self._file_size + self._zinfo.CRC = self._crc + self._zinfo.file_size = self._file_size - # Write updated header info - if self._zinfo.flag_bits & 0x08: - # Write CRC and file sizes after the file data - fmt = ' ZIP64_LIMIT: - raise RuntimeError('File size unexpectedly exceeded ZIP64 ' - 'limit') - if self._compress_size > ZIP64_LIMIT: - raise RuntimeError('Compressed size unexpectedly exceeded ' - 'ZIP64 limit') - # Seek backwards and write file header (which will now include - # correct CRC and file sizes) + # Write updated header info + if self._zinfo.flag_bits & 0x08: + # Write CRC and file sizes after the file data + fmt = ' ZIP64_LIMIT: + raise RuntimeError( + 'File size unexpectedly exceeded ZIP64 limit') + if self._compress_size > ZIP64_LIMIT: + raise RuntimeError( + 'Compressed size unexpectedly exceeded ZIP64 limit') + # Seek backwards and write file header (which will now include + # correct CRC and file sizes) - # Preserve current position in file - self._zipfile.start_dir = self._fileobj.tell() - self._fileobj.seek(self._zinfo.header_offset) - self._fileobj.write(self._zinfo.FileHeader(self._zip64)) - self._fileobj.seek(self._zipfile.start_dir) + # Preserve current position in file + self._zipfile.start_dir = self._fileobj.tell() + self._fileobj.seek(self._zinfo.header_offset) + self._fileobj.write(self._zinfo.FileHeader(self._zip64)) + self._fileobj.seek(self._zipfile.start_dir) + + # Successfully written: Add file to our caches + self._zipfile.filelist.append(self._zinfo) + self._zipfile.NameToInfo[self._zinfo.filename] = self._zinfo + finally: + self._zipfile._writing = False - self._zipfile._writing = False - # Successfully written: Add file to our caches - self._zipfile.filelist.append(self._zinfo) - self._zipfile.NameToInfo[self._zinfo.filename] = self._zinfo class ZipFile: """ Class with methods to open, read, write, close, list zip files. diff --git a/Misc/NEWS.d/next/Library/2019-03-26-14-20-59.bpo-36434.PTdidw.rst b/Misc/NEWS.d/next/Library/2019-03-26-14-20-59.bpo-36434.PTdidw.rst new file mode 100644 index 00000000000..6e3e050c89b --- /dev/null +++ b/Misc/NEWS.d/next/Library/2019-03-26-14-20-59.bpo-36434.PTdidw.rst @@ -0,0 +1 @@ +Errors during writing to a ZIP file no longer prevent to properly close it.