Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 37 additions & 0 deletions Lib/test/test_zipfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -402,6 +402,43 @@ def test_per_file_compresslevel(self):
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)
Expand Down
85 changes: 44 additions & 41 deletions Lib/zipfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -1101,47 +1101,50 @@ def write(self, data):
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

# 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)

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
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)

# 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



class ZipFile:
""" Class with methods to open, read, write, close, list zip files.
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Errors during writing to a ZIP file no longer prevent to properly close it.