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
GitHub
parent
7a465cb5ee
commit
2524fdefc9
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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 = '<LLQQ' if self._zip64 else '<LLLL'
|
||||
self._fileobj.write(struct.pack(fmt, _DD_SIGNATURE, self._zinfo.CRC,
|
||||
self._zinfo.compress_size, self._zinfo.file_size))
|
||||
self._zipfile.start_dir = self._fileobj.tell()
|
||||
else:
|
||||
if not self._zip64:
|
||||
if self._file_size > 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 = '<LLQQ' if self._zip64 else '<LLLL'
|
||||
self._fileobj.write(struct.pack(fmt, _DD_SIGNATURE, self._zinfo.CRC,
|
||||
self._zinfo.compress_size, self._zinfo.file_size))
|
||||
self._zipfile.start_dir = self._fileobj.tell()
|
||||
else:
|
||||
if not self._zip64:
|
||||
if self._file_size > 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.
|
||||
|
|
|
|||
|
|
@ -0,0 +1 @@
|
|||
Errors during writing to a ZIP file no longer prevent to properly close it.
|
||||
Loading…
Reference in New Issue