Added the --compression option to bdist_wheel (#316)

Fixes #313.

4 files changed, 38 insertions, 5 deletions

diff --git a/docs/news.rst b/docs/news.rst
index 8939c15..2911fa6 100644
--- a/docs/news.rst
+++ b/docs/news.rst
@@ -6,6 +6,7 @@ Release Notes
 - Dropped Python 3.4 support
 - Added automatic platform tag detection for macOS binary wheels
   (PR by Grzegorz Bokota)
+- Added the ``--compression=`` option to the ``bdist_wheel`` command
 - Fixed PyPy tag generation to work with the updated semantics (#328)
 - Updated project packaging and testing configuration for :pep:`517`
 - Moved the contents of setup.py to setup.cfg
diff --git a/src/wheel/bdist_wheel.py b/src/wheel/bdist_wheel.py
index ce82eb6..5cece1f 100644
--- a/src/wheel/bdist_wheel.py
+++ b/src/wheel/bdist_wheel.py
@@ -9,6 +9,7 @@ import shutil
 import stat
 import sys
 import re
+from collections import OrderedDict
 from email.generator import Generator
 from distutils.core import Command
 from distutils.sysconfig import get_python_version
@@ -16,6 +17,7 @@ from distutils import log as logger
 from glob import iglob
 from shutil import rmtree
 from warnings import warn
+from zipfile import ZIP_DEFLATED, ZIP_STORED
 
 import pkg_resources
 
@@ -51,6 +53,11 @@ class bdist_wheel(Command):
 
     description = 'create a wheel distribution'
 
+    supported_compressions = OrderedDict([
+        ('stored', ZIP_STORED),
+        ('deflated', ZIP_DEFLATED)
+    ])
+
     user_options = [('bdist-dir=', 'b',
                      "temporary directory for creating the distribution"),
                     ('plat-name=', 'p',
@@ -75,6 +82,10 @@ class bdist_wheel(Command):
                     ('universal', None,
                      "make a universal wheel"
                      " (default: false)"),
+                    ('compression=', None,
+                     "zipfile compression (one of: {})"
+                     " (default: 'deflated')"
+                     .format(', '.join(supported_compressions))),
                     ('python-tag=', None,
                      "Python implementation compatibility tag"
                      " (default: py%s)" % get_impl_ver()[0]),
@@ -104,6 +115,7 @@ class bdist_wheel(Command):
         self.owner = None
         self.group = None
         self.universal = False
+        self.compression = 'deflated'
         self.python_tag = 'py' + get_impl_ver()[0]
         self.build_number = None
         self.py_limited_api = False
@@ -117,6 +129,11 @@ class bdist_wheel(Command):
         self.data_dir = self.wheel_dist_name + '.data'
         self.plat_name_supplied = self.plat_name is not None
 
+        try:
+            self.compression = self.supported_compressions[self.compression]
+        except KeyError:
+            raise ValueError('Unsupported compression: {}'.format(self.compression))
+
         need_options = ('dist_dir', 'plat_name', 'skip_build')
 
         self.set_undefined_options('bdist',
@@ -264,7 +281,7 @@ class bdist_wheel(Command):
             os.makedirs(self.dist_dir)
 
         wheel_path = os.path.join(self.dist_dir, archive_basename + '.whl')
-        with WheelFile(wheel_path, 'w') as wf:
+        with WheelFile(wheel_path, 'w', self.compression) as wf:
             wf.write_files(archive_root)
 
         # Add to 'Distribution.dist_files' so that the "upload" command works
diff --git a/src/wheel/wheelfile.py b/src/wheel/wheelfile.py
index fa00258..acc5dab 100644
--- a/src/wheel/wheelfile.py
+++ b/src/wheel/wheelfile.py
@@ -35,13 +35,13 @@ class WheelFile(ZipFile):
 
     _default_algorithm = hashlib.sha256
 
-    def __init__(self, file, mode='r'):
+    def __init__(self, file, mode='r', compression=ZIP_DEFLATED):
         basename = os.path.basename(file)
         self.parsed_filename = WHEEL_INFO_RE.match(basename)
         if not basename.endswith('.whl') or self.parsed_filename is None:
             raise WheelError("Bad wheel filename {!r}".format(basename))
 
-        ZipFile.__init__(self, file, mode, compression=ZIP_DEFLATED, allowZip64=True)
+        ZipFile.__init__(self, file, mode, compression=compression, allowZip64=True)
 
         self.dist_info_path = '{}.dist-info'.format(self.parsed_filename.group('namever'))
         self.record_path = self.dist_info_path + '/RECORD'
@@ -134,7 +134,7 @@ class WheelFile(ZipFile):
 
         zinfo = ZipInfo(arcname or filename, date_time=get_zipinfo_datetime(st.st_mtime))
         zinfo.external_attr = (stat.S_IMODE(st.st_mode) | stat.S_IFMT(st.st_mode)) << 16
-        zinfo.compress_type = ZIP_DEFLATED
+        zinfo.compress_type = compress_type or self.compression
         self.writestr(zinfo, data, compress_type)
 
     def writestr(self, zinfo_or_arcname, bytes, compress_type=None):
@@ -162,7 +162,7 @@ class WheelFile(ZipFile):
             ))
             writer.writerow((format(self.record_path), "", ""))
             zinfo = ZipInfo(native(self.record_path), date_time=get_zipinfo_datetime())
-            zinfo.compress_type = ZIP_DEFLATED
+            zinfo.compress_type = self.compression
             zinfo.external_attr = 0o664 << 16
             self.writestr(zinfo, as_bytes(data.getvalue()))
 
diff --git a/tests/test_bdist_wheel.py b/tests/test_bdist_wheel.py
index de8d46f..67df435 100644
--- a/tests/test_bdist_wheel.py
+++ b/tests/test_bdist_wheel.py
@@ -8,6 +8,7 @@ from zipfile import ZipFile
 
 import pytest
 
+from wheel.bdist_wheel import bdist_wheel
 from wheel.wheelfile import WheelFile
 
 DEFAULT_FILES = {
@@ -128,3 +129,17 @@ def test_build_from_readonly_tree(dummy_dist, monkeypatch, tmpdir):
             os.chmod(os.path.join(root, fname), stat.S_IREAD)
 
     subprocess.check_call([sys.executable, 'setup.py', 'bdist_wheel'])
+
+
+@pytest.mark.parametrize('option, compress_type', list(bdist_wheel.supported_compressions.items()),
+                         ids=list(bdist_wheel.supported_compressions))
+def test_compression(dummy_dist, monkeypatch, tmpdir, option, compress_type):
+    monkeypatch.chdir(dummy_dist)
+    subprocess.check_call([sys.executable, 'setup.py', 'bdist_wheel', '-b', str(tmpdir),
+                           '--universal', '--compression={}'.format(option)])
+    with WheelFile('dist/dummy_dist-1.0-py2.py3-none-any.whl') as wf:
+        filenames = set(wf.namelist())
+        assert 'dummy_dist-1.0.dist-info/RECORD' in filenames
+        assert 'dummy_dist-1.0.dist-info/METADATA' in filenames
+        for zinfo in wf.filelist:
+            assert zinfo.compress_type == compress_type