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
diff --git a/Lib/zipfile.py b/Lib/zipfile.py
index 37ce328..f9db45f 100644
--- a/Lib/zipfile.py
+++ b/Lib/zipfile.py
@@ -295,6 +295,7 @@
'filename',
'date_time',
'compress_type',
+ '_compresslevel',
'comment',
'extra',
'create_system',
@@ -334,6 +335,7 @@
# Standard values:
self.compress_type = ZIP_STORED # Type of compression for the file
+ self._compresslevel = None # Level for the compressor
self.comment = b"" # Comment for each file
self.extra = b"" # ZIP extra data
if sys.platform == 'win32':
@@ -654,12 +656,16 @@
raise NotImplementedError("That compression method is not supported")
-def _get_compressor(compress_type):
+def _get_compressor(compress_type, compresslevel=None):
if compress_type == ZIP_DEFLATED:
- return zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
- zlib.DEFLATED, -15)
+ if compresslevel is not None:
+ return zlib.compressobj(compresslevel, zlib.DEFLATED, -15)
+ return zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -15)
elif compress_type == ZIP_BZIP2:
+ if compresslevel is not None:
+ return bz2.BZ2Compressor(compresslevel)
return bz2.BZ2Compressor()
+ # compresslevel is ignored for ZIP_LZMA
elif compress_type == ZIP_LZMA:
return LZMACompressor()
else:
@@ -963,7 +969,8 @@
self._zinfo = zinfo
self._zip64 = zip64
self._zipfile = zf
- self._compressor = _get_compressor(zinfo.compress_type)
+ self._compressor = _get_compressor(zinfo.compress_type,
+ zinfo._compresslevel)
self._file_size = 0
self._compress_size = 0
self._crc = 0
@@ -1035,7 +1042,8 @@
class ZipFile:
""" Class with methods to open, read, write, close, list zip files.
- z = ZipFile(file, mode="r", compression=ZIP_STORED, allowZip64=True)
+ z = ZipFile(file, mode="r", compression=ZIP_STORED, allowZip64=True,
+ compresslevel=None)
file: Either the path to the file, or a file-like object.
If it is a path, the file will be opened and closed by ZipFile.
@@ -1046,13 +1054,19 @@
allowZip64: if True ZipFile will create files with ZIP64 extensions when
needed, otherwise it will raise an exception when this would
be necessary.
+ compresslevel: None (default for the given compression type) or an integer
+ specifying the level to pass to the compressor.
+ When using ZIP_STORED or ZIP_LZMA this keyword has no effect.
+ When using ZIP_DEFLATED integers 0 through 9 are accepted.
+ When using ZIP_BZIP2 integers 1 through 9 are accepted.
"""
fp = None # Set here since __del__ checks it
_windows_illegal_name_trans_table = None
- def __init__(self, file, mode="r", compression=ZIP_STORED, allowZip64=True):
+ def __init__(self, file, mode="r", compression=ZIP_STORED, allowZip64=True,
+ compresslevel=None):
"""Open the ZIP file with mode read 'r', write 'w', exclusive create 'x',
or append 'a'."""
if mode not in ('r', 'w', 'x', 'a'):
@@ -1066,6 +1080,7 @@
self.NameToInfo = {} # Find file info given name
self.filelist = [] # List of ZipInfo instances for archive
self.compression = compression # Method of compression
+ self.compresslevel = compresslevel
self.mode = mode
self.pwd = None
self._comment = b''
@@ -1342,6 +1357,7 @@
elif mode == 'w':
zinfo = ZipInfo(name)
zinfo.compress_type = self.compression
+ zinfo._compresslevel = self.compresslevel
else:
# Get info object for name
zinfo = self.getinfo(name)
@@ -1575,7 +1591,8 @@
raise LargeZipFile(requires_zip64 +
" would require ZIP64 extensions")
- def write(self, filename, arcname=None, compress_type=None):
+ def write(self, filename, arcname=None,
+ compress_type=None, compresslevel=None):
"""Put the bytes from filename into the archive under the name
arcname."""
if not self.fp:
@@ -1597,6 +1614,11 @@
else:
zinfo.compress_type = self.compression
+ if compresslevel is not None:
+ zinfo._compresslevel = compresslevel
+ else:
+ zinfo._compresslevel = self.compresslevel
+
if zinfo.is_dir():
with self._lock:
if self._seekable:
@@ -1617,7 +1639,8 @@
with open(filename, "rb") as src, self.open(zinfo, 'w') as dest:
shutil.copyfileobj(src, dest, 1024*8)
- def writestr(self, zinfo_or_arcname, data, compress_type=None):
+ def writestr(self, zinfo_or_arcname, data,
+ compress_type=None, compresslevel=None):
"""Write a file into the archive. The contents is 'data', which
may be either a 'str' or a 'bytes' instance; if it is a 'str',
it is encoded as UTF-8 first.
@@ -1629,6 +1652,7 @@
zinfo = ZipInfo(filename=zinfo_or_arcname,
date_time=time.localtime(time.time())[:6])
zinfo.compress_type = self.compression
+ zinfo._compresslevel = self.compresslevel
if zinfo.filename[-1] == '/':
zinfo.external_attr = 0o40775 << 16 # drwxrwxr-x
zinfo.external_attr |= 0x10 # MS-DOS directory flag
@@ -1648,6 +1672,9 @@
if compress_type is not None:
zinfo.compress_type = compress_type
+ if compresslevel is not None:
+ zinfo._compresslevel = compresslevel
+
zinfo.file_size = len(data) # Uncompressed size
with self._lock:
with self.open(zinfo, mode='w') as dest: