bpo-21417: Add compresslevel= to the zipfile module (GH-5385)

This allows the compression level to be specified when writing zipfiles
(for the entire file *and* overridden on a per-file basis).

Contributed by Bo Bayles
This commit is contained in:
Bo Bayles 2018-01-29 23:54:07 -06:00 committed by Gregory P. Smith
parent f4d644f36f
commit ce237c7d58
4 changed files with 104 additions and 24 deletions

View file

@ -53,9 +53,10 @@ class AbstractTestsWithSourceFile:
with open(TESTFN, "wb") as fp:
fp.write(self.data)
def make_test_archive(self, f, compression):
def make_test_archive(self, f, compression, compresslevel=None):
kwargs = {'compression': compression, 'compresslevel': compresslevel}
# Create the ZIP archive
with zipfile.ZipFile(f, "w", compression) as zipfp:
with zipfile.ZipFile(f, "w", **kwargs) as zipfp:
zipfp.write(TESTFN, "another.name")
zipfp.write(TESTFN, TESTFN)
zipfp.writestr("strfile", self.data)
@ -63,8 +64,8 @@ class AbstractTestsWithSourceFile:
for line in self.line_gen:
f.write(line)
def zip_test(self, f, compression):
self.make_test_archive(f, compression)
def zip_test(self, f, compression, compresslevel=None):
self.make_test_archive(f, compression, compresslevel)
# Read the ZIP archive
with zipfile.ZipFile(f, "r", compression) as zipfp:
@ -297,6 +298,22 @@ class AbstractTestsWithSourceFile:
info = zipfp.getinfo('b.txt')
self.assertEqual(info.compress_type, self.compression)
def test_writestr_compresslevel(self):
zipfp = zipfile.ZipFile(TESTFN2, "w", compresslevel=1)
zipfp.writestr("a.txt", "hello world", compress_type=self.compression)
zipfp.writestr("b.txt", "hello world", compress_type=self.compression,
compresslevel=2)
# Compression level follows the constructor.
a_info = zipfp.getinfo('a.txt')
self.assertEqual(a_info.compress_type, self.compression)
self.assertEqual(a_info._compresslevel, 1)
# Compression level is overridden.
b_info = zipfp.getinfo('b.txt')
self.assertEqual(b_info.compress_type, self.compression)
self.assertEqual(b_info._compresslevel, 2)
def test_read_return_size(self):
# Issue #9837: ZipExtFile.read() shouldn't return more bytes
# than requested.
@ -370,6 +387,21 @@ class AbstractTestsWithSourceFile:
self.assertIn('[closed]', repr(zipopen))
self.assertIn('[closed]', repr(zipfp))
def test_compresslevel_basic(self):
for f in get_files(self):
self.zip_test(f, self.compression, compresslevel=9)
def test_per_file_compresslevel(self):
"""Check that files within a Zip archive can have different
compression levels."""
with zipfile.ZipFile(TESTFN2, "w", compresslevel=1) as zipfp:
zipfp.write(TESTFN, 'compress_1')
zipfp.write(TESTFN, 'compress_9', compresslevel=9)
one_info = zipfp.getinfo('compress_1')
nine_info = zipfp.getinfo('compress_9')
self.assertEqual(one_info._compresslevel, 1)
self.assertEqual(nine_info._compresslevel, 9)
def tearDown(self):
unlink(TESTFN)
unlink(TESTFN2)