bpo-36434: Properly handle writing errors in ZIP files. (GH-12559)

Errors during writing no longer prevent to properly close
the ZIP file.
This commit is contained in:
Serhiy Storchaka 2019-03-30 08:25:19 +02:00 committed by GitHub
parent 7a465cb5ee
commit 2524fdefc9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 78 additions and 37 deletions

View File

@ -402,6 +402,43 @@ class AbstractTestsWithSourceFile:
self.assertEqual(one_info._compresslevel, 1) self.assertEqual(one_info._compresslevel, 1)
self.assertEqual(nine_info._compresslevel, 9) 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): def tearDown(self):
unlink(TESTFN) unlink(TESTFN)
unlink(TESTFN2) unlink(TESTFN2)

View File

@ -1105,47 +1105,50 @@ class _ZipWriteFile(io.BufferedIOBase):
def close(self): def close(self):
if self.closed: if self.closed:
return return
super().close() try:
# Flush any data from the compressor, and update header info super().close()
if self._compressor: # Flush any data from the compressor, and update header info
buf = self._compressor.flush() if self._compressor:
self._compress_size += len(buf) buf = self._compressor.flush()
self._fileobj.write(buf) self._compress_size += len(buf)
self._zinfo.compress_size = self._compress_size self._fileobj.write(buf)
else: self._zinfo.compress_size = self._compress_size
self._zinfo.compress_size = self._file_size else:
self._zinfo.CRC = self._crc self._zinfo.compress_size = self._file_size
self._zinfo.file_size = self._file_size self._zinfo.CRC = self._crc
self._zinfo.file_size = self._file_size
# Write updated header info # Write updated header info
if self._zinfo.flag_bits & 0x08: if self._zinfo.flag_bits & 0x08:
# Write CRC and file sizes after the file data # Write CRC and file sizes after the file data
fmt = '<LLQQ' if self._zip64 else '<LLLL' fmt = '<LLQQ' if self._zip64 else '<LLLL'
self._fileobj.write(struct.pack(fmt, _DD_SIGNATURE, self._zinfo.CRC, self._fileobj.write(struct.pack(fmt, _DD_SIGNATURE, self._zinfo.CRC,
self._zinfo.compress_size, self._zinfo.file_size)) self._zinfo.compress_size, self._zinfo.file_size))
self._zipfile.start_dir = self._fileobj.tell() self._zipfile.start_dir = self._fileobj.tell()
else: else:
if not self._zip64: if not self._zip64:
if self._file_size > ZIP64_LIMIT: if self._file_size > ZIP64_LIMIT:
raise RuntimeError('File size unexpectedly exceeded ZIP64 ' raise RuntimeError(
'limit') 'File size unexpectedly exceeded ZIP64 limit')
if self._compress_size > ZIP64_LIMIT: if self._compress_size > ZIP64_LIMIT:
raise RuntimeError('Compressed size unexpectedly exceeded ' raise RuntimeError(
'ZIP64 limit') 'Compressed size unexpectedly exceeded ZIP64 limit')
# Seek backwards and write file header (which will now include # Seek backwards and write file header (which will now include
# correct CRC and file sizes) # correct CRC and file sizes)
# Preserve current position in file # Preserve current position in file
self._zipfile.start_dir = self._fileobj.tell() self._zipfile.start_dir = self._fileobj.tell()
self._fileobj.seek(self._zinfo.header_offset) self._fileobj.seek(self._zinfo.header_offset)
self._fileobj.write(self._zinfo.FileHeader(self._zip64)) self._fileobj.write(self._zinfo.FileHeader(self._zip64))
self._fileobj.seek(self._zipfile.start_dir) 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 ZipFile:
""" Class with methods to open, read, write, close, list zip files. """ Class with methods to open, read, write, close, list zip files.

View File

@ -0,0 +1 @@
Errors during writing to a ZIP file no longer prevent to properly close it.