Skip to content

Commit 4724ba9

Browse files
bpo-36434: Properly handle writing errors in ZIP files. (GH-12559) (GH-12628)
Errors during writing no longer prevent to properly close the ZIP file. (cherry picked from commit 2524fde) Co-authored-by: Serhiy Storchaka <storchaka@gmail.com>
1 parent 128e40f commit 4724ba9

3 files changed

Lines changed: 82 additions & 41 deletions

File tree

Lib/test/test_zipfile.py

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -402,6 +402,43 @@ def test_per_file_compresslevel(self):
402402
self.assertEqual(one_info._compresslevel, 1)
403403
self.assertEqual(nine_info._compresslevel, 9)
404404

405+
def test_writing_errors(self):
406+
class BrokenFile(io.BytesIO):
407+
def write(self, data):
408+
nonlocal count
409+
if count is not None:
410+
if count == stop:
411+
raise OSError
412+
count += 1
413+
super().write(data)
414+
415+
stop = 0
416+
while True:
417+
testfile = BrokenFile()
418+
count = None
419+
with zipfile.ZipFile(testfile, 'w', self.compression) as zipfp:
420+
with zipfp.open('file1', 'w') as f:
421+
f.write(b'data1')
422+
count = 0
423+
try:
424+
with zipfp.open('file2', 'w') as f:
425+
f.write(b'data2')
426+
except OSError:
427+
stop += 1
428+
else:
429+
break
430+
finally:
431+
count = None
432+
with zipfile.ZipFile(io.BytesIO(testfile.getvalue())) as zipfp:
433+
self.assertEqual(zipfp.namelist(), ['file1'])
434+
self.assertEqual(zipfp.read('file1'), b'data1')
435+
436+
with zipfile.ZipFile(io.BytesIO(testfile.getvalue())) as zipfp:
437+
self.assertEqual(zipfp.namelist(), ['file1', 'file2'])
438+
self.assertEqual(zipfp.read('file1'), b'data1')
439+
self.assertEqual(zipfp.read('file2'), b'data2')
440+
441+
405442
def tearDown(self):
406443
unlink(TESTFN)
407444
unlink(TESTFN2)

Lib/zipfile.py

Lines changed: 44 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -1101,47 +1101,50 @@ def write(self, data):
11011101
def close(self):
11021102
if self.closed:
11031103
return
1104-
super().close()
1105-
# Flush any data from the compressor, and update header info
1106-
if self._compressor:
1107-
buf = self._compressor.flush()
1108-
self._compress_size += len(buf)
1109-
self._fileobj.write(buf)
1110-
self._zinfo.compress_size = self._compress_size
1111-
else:
1112-
self._zinfo.compress_size = self._file_size
1113-
self._zinfo.CRC = self._crc
1114-
self._zinfo.file_size = self._file_size
1115-
1116-
# Write updated header info
1117-
if self._zinfo.flag_bits & 0x08:
1118-
# Write CRC and file sizes after the file data
1119-
fmt = '<LLQQ' if self._zip64 else '<LLLL'
1120-
self._fileobj.write(struct.pack(fmt, _DD_SIGNATURE, self._zinfo.CRC,
1121-
self._zinfo.compress_size, self._zinfo.file_size))
1122-
self._zipfile.start_dir = self._fileobj.tell()
1123-
else:
1124-
if not self._zip64:
1125-
if self._file_size > ZIP64_LIMIT:
1126-
raise RuntimeError('File size unexpectedly exceeded ZIP64 '
1127-
'limit')
1128-
if self._compress_size > ZIP64_LIMIT:
1129-
raise RuntimeError('Compressed size unexpectedly exceeded '
1130-
'ZIP64 limit')
1131-
# Seek backwards and write file header (which will now include
1132-
# correct CRC and file sizes)
1133-
1134-
# Preserve current position in file
1135-
self._zipfile.start_dir = self._fileobj.tell()
1136-
self._fileobj.seek(self._zinfo.header_offset)
1137-
self._fileobj.write(self._zinfo.FileHeader(self._zip64))
1138-
self._fileobj.seek(self._zipfile.start_dir)
1139-
1140-
self._zipfile._writing = False
1141-
1142-
# Successfully written: Add file to our caches
1143-
self._zipfile.filelist.append(self._zinfo)
1144-
self._zipfile.NameToInfo[self._zinfo.filename] = self._zinfo
1104+
try:
1105+
super().close()
1106+
# Flush any data from the compressor, and update header info
1107+
if self._compressor:
1108+
buf = self._compressor.flush()
1109+
self._compress_size += len(buf)
1110+
self._fileobj.write(buf)
1111+
self._zinfo.compress_size = self._compress_size
1112+
else:
1113+
self._zinfo.compress_size = self._file_size
1114+
self._zinfo.CRC = self._crc
1115+
self._zinfo.file_size = self._file_size
1116+
1117+
# Write updated header info
1118+
if self._zinfo.flag_bits & 0x08:
1119+
# Write CRC and file sizes after the file data
1120+
fmt = '<LLQQ' if self._zip64 else '<LLLL'
1121+
self._fileobj.write(struct.pack(fmt, _DD_SIGNATURE, self._zinfo.CRC,
1122+
self._zinfo.compress_size, self._zinfo.file_size))
1123+
self._zipfile.start_dir = self._fileobj.tell()
1124+
else:
1125+
if not self._zip64:
1126+
if self._file_size > ZIP64_LIMIT:
1127+
raise RuntimeError(
1128+
'File size unexpectedly exceeded ZIP64 limit')
1129+
if self._compress_size > ZIP64_LIMIT:
1130+
raise RuntimeError(
1131+
'Compressed size unexpectedly exceeded ZIP64 limit')
1132+
# Seek backwards and write file header (which will now include
1133+
# correct CRC and file sizes)
1134+
1135+
# Preserve current position in file
1136+
self._zipfile.start_dir = self._fileobj.tell()
1137+
self._fileobj.seek(self._zinfo.header_offset)
1138+
self._fileobj.write(self._zinfo.FileHeader(self._zip64))
1139+
self._fileobj.seek(self._zipfile.start_dir)
1140+
1141+
# Successfully written: Add file to our caches
1142+
self._zipfile.filelist.append(self._zinfo)
1143+
self._zipfile.NameToInfo[self._zinfo.filename] = self._zinfo
1144+
finally:
1145+
self._zipfile._writing = False
1146+
1147+
11451148

11461149
class ZipFile:
11471150
""" Class with methods to open, read, write, close, list zip files.
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Errors during writing to a ZIP file no longer prevent to properly close it.

0 commit comments

Comments
 (0)