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:
parent
7a465cb5ee
commit
2524fdefc9
|
@ -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)
|
||||||
|
|
|
@ -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.
|
||||||
|
|
|
@ -0,0 +1 @@
|
||||||
|
Errors during writing to a ZIP file no longer prevent to properly close it.
|
Loading…
Reference in New Issue