WheelFile: fix default compressargs in write_*

It was taking `ZipFile` defaults instead of the `ZIP_DEFLATED` / Those set in `__init__`.
MrMino · Aug 2, 2021 · 5262b8c · 5262b8c
1 parent 57df729
commit 5262b8c
Show file tree

Hide file tree

Showing 2 changed files with 145 additions and 0 deletions.
diff --git a/tests/test_wheelfile.py b/tests/test_wheelfile.py
@@ -945,3 +945,140 @@ def test_write_distinfo_sets_the_right_compresslevel(self, wf, tmp_file):
         wf.write_distinfo(tmp_file, compresslevel=7)
         arcpath = wf.distinfo_dirname + '/' + tmp_file.name
         assert wf.zipfile.getinfo(arcpath)._compresslevel == 7
+
+    def test_write_default_compress_type_is_deflate(self, buf, tmp_file):
+        wf = WheelFile(buf, 'w', distname='_', version='0')
+        wf.write(tmp_file)
+        assert wf.infolist()[0].compress_type == ZIP_DEFLATED
+
+    def test_write_data_default_compress_type_is_deflate(self, buf, tmp_file):
+        wf = WheelFile(buf, 'w', distname='_', version='0')
+        wf.write_data(tmp_file, 'section')
+        assert wf.infolist()[0].compress_type == ZIP_DEFLATED
+
+    def test_write_distinfo_default_compress_type_is_deflate(self, buf,
+                                                             tmp_file):
+        wf = WheelFile(buf, 'w', distname='_', version='0')
+        wf.write_distinfo(tmp_file)
+        assert wf.infolist()[0].compress_type == ZIP_DEFLATED
+
+    def test_write_default_compress_type_is_from_init(self, buf, tmp_file):
+        wf = WheelFile(buf, 'w', distname='_', version='0',
+                       compression=ZIP_BZIP2)
+        wf.write(tmp_file)
+        assert wf.infolist()[0].compress_type == ZIP_BZIP2
+
+    def test_write_data_default_compress_type_is_from_init(self, buf, tmp_file):
+        wf = WheelFile(buf, 'w', distname='_', version='0',
+                       compression=ZIP_BZIP2)
+        wf.write_data(tmp_file, 'section')
+        assert wf.infolist()[0].compress_type == ZIP_BZIP2
+
+    def test_write_distinfo_default_compress_type_is_from_init(self, buf,
+                                                               tmp_file):
+        wf = WheelFile(buf, 'w', distname='_', version='0',
+                       compression=ZIP_BZIP2)
+        wf.write_distinfo(tmp_file)
+        assert wf.infolist()[0].compress_type == ZIP_BZIP2
+
+    def test_write_default_compresslevel_is_none(self, buf, tmp_file):
+        wf = WheelFile(buf, 'w', distname='_', version='0')
+        wf.write(tmp_file)
+        assert wf.infolist()[0]._compresslevel is None
+
+    def test_write_data_default_compresslevel_is_none(self, buf, tmp_file):
+        wf = WheelFile(buf, 'w', distname='_', version='0')
+        wf.write_data(tmp_file, 'section')
+        assert wf.infolist()[0]._compresslevel is None
+
+    def test_write_distinfo_default_compresslevel_is_none(self, buf, tmp_file):
+        wf = WheelFile(buf, 'w', distname='_', version='0')
+        wf.write_distinfo(tmp_file)
+        assert wf.infolist()[0]._compresslevel is None
+
+    def test_write_default_compresslevel_is_from_init(self, buf,
+                                                      tmp_file):
+        wf = WheelFile(buf, 'w', distname='_', version='0',
+                       compresslevel=9)
+        wf.write(tmp_file)
+        assert wf.infolist()[0]._compresslevel == 9
+
+    def test_write_data_default_compresslevel_is_from_init(self, buf,
+                                                           tmp_file):
+        wf = WheelFile(buf, 'w', distname='_', version='0',
+                       compresslevel=9)
+        wf.write_data(tmp_file, 'section')
+        assert wf.infolist()[0]._compresslevel == 9
+
+    def test_write_distinfo_default_compresslevel_is_from_init(self, buf,
+                                                               tmp_file):
+        wf = WheelFile(buf, 'w', distname='_', version='0',
+                       compresslevel=9)
+        wf.write_distinfo(tmp_file)
+        assert wf.infolist()[0]._compresslevel == 9
+
+    def test_writestr_default_compress_type_is_deflate(self, buf):
+        wf = WheelFile(buf, 'w', distname='_', version='0')
+        wf.writestr('file', b'data')
+        assert wf.infolist()[0].compress_type == ZIP_DEFLATED
+
+    def test_writestr_data_default_compress_type_is_deflate(self, buf):
+        wf = WheelFile(buf, 'w', distname='_', version='0')
+        wf.writestr_data('section', 'file', b'data')
+        assert wf.infolist()[0].compress_type == ZIP_DEFLATED
+
+    def test_writestr_distinfo_default_compress_type_is_deflate(self, buf):
+        wf = WheelFile(buf, 'w', distname='_', version='0')
+        wf.writestr_distinfo('file', b'data')
+        assert wf.infolist()[0].compress_type == ZIP_DEFLATED
+
+    def test_writestr_default_compress_type_is_from_init(self, buf):
+        wf = WheelFile(buf, 'w', distname='_', version='0',
+                       compression=ZIP_BZIP2)
+        wf.writestr('file', b'data')
+        assert wf.infolist()[0].compress_type == ZIP_BZIP2
+
+    def test_writestr_data_default_compress_type_is_from_init(self, buf):
+        wf = WheelFile(buf, 'w', distname='_', version='0',
+                       compression=ZIP_BZIP2)
+        wf.writestr_data('section', 'file', b'data')
+        assert wf.infolist()[0].compress_type == ZIP_BZIP2
+
+    def test_writestr_distinfo_default_compress_type_is_from_init(self, buf):
+        wf = WheelFile(buf, 'w', distname='_', version='0',
+                       compression=ZIP_BZIP2)
+        wf.writestr_distinfo('file', b'data')
+        assert wf.infolist()[0].compress_type == ZIP_BZIP2
+
+    def test_writestr_default_compresslevel_is_none(self, buf):
+        wf = WheelFile(buf, 'w', distname='_', version='0')
+        wf.writestr('file', b'data')
+        assert wf.infolist()[0]._compresslevel is None
+
+    def test_writestr_data_default_compresslevel_is_none(self, buf):
+        wf = WheelFile(buf, 'w', distname='_', version='0')
+        wf.writestr_data('section', 'file', b'data')
+        assert wf.infolist()[0]._compresslevel is None
+
+    def test_writestr_distinfo_default_compresslevel_is_none(self, buf):
+        wf = WheelFile(buf, 'w', distname='_', version='0')
+        wf.writestr_distinfo('file', b'data')
+        assert wf.infolist()[0]._compresslevel is None
+
+    def test_writestr_default_compresslevel_is_from_init(self, buf):
+        wf = WheelFile(buf, 'w', distname='_', version='0',
+                       compresslevel=9)
+        wf.writestr('file', b'data')
+        assert wf.infolist()[0]._compresslevel == 9
+
+    def test_writestr_data_default_compresslevel_is_from_init(self, buf):
+        wf = WheelFile(buf, 'w', distname='_', version='0',
+                       compresslevel=9)
+        wf.writestr_data('section', 'file', b'data')
+        assert wf.infolist()[0]._compresslevel == 9
+
+    def test_writestr_distinfo_default_compresslevel_is_from_init(self, buf):
+        wf = WheelFile(buf, 'w', distname='_', version='0',
+                       compresslevel=9)
+        wf.writestr_distinfo('file', b'data')
+        assert wf.infolist()[0]._compresslevel == 9
diff --git a/wheelfile.py b/wheelfile.py
@@ -1942,6 +1942,14 @@ def _write_to_zip(self, filename, arcname, skipdir, compress_type,
         zinfo = zipfile.ZipInfo.from_file(
             filename, arcname, strict_timestamps=self._strict_timestamps
         )
+
+        # Since we construct ZipInfo manually here, we have to propagate
+        # defaults ourselves.
+        if compress_type is None:
+            compress_type = self.zipfile.compression
+        if compresslevel is None:
+            compresslevel = self.zipfile.compresslevel
+
         if zinfo.is_dir():
             if skipdir:
                 return