Support repack(removed)

Author: danny0838

Doc/library/zipfile.rst

@@ -527,20 +527,33 @@ ZipFile Objects
    a path is provided.
 
    This does not physically remove the local file entry from the archive;
-   the ZIP file size remains unchanged. Use :meth:`ZipFile.repack` afterwards
+   the ZIP file size remains unchanged. Call :meth:`ZipFile.repack` afterwards
    to reclaim space.
 
    The archive must be opened with mode ``'w'``, ``'x'`` or ``'a'``.
 
+   Returns the removed :class:`ZipInfo` instance.
+
    Calling :meth:`remove` on a closed ZipFile will raise a :exc:`ValueError`.
 
    .. versionadded:: next
 
 
-.. method:: ZipFile.repack(*, strict_descriptor=False[, chunk_size])
+.. method:: ZipFile.repack(removed=None, *, \
+                           strict_descriptor=False[, chunk_size])
+
+   Rewrites the archive to remove stale local file entries, shrinking the ZIP
+   file size.
+
+   If *removed* is provided, it must be a sequence of :class:`ZipInfo` objects
+   representing removed entries; only their corresponding local file entries
+   will be removed.
 
-   Rewrites the archive to remove local file entries that are no longer
-   referenced, shrinking the ZIP file size.
+   If *removed* is not provided, local file entries no longer referenced in the
+   central directory will be removed. The algorithm assumes that local file
+   entries are stored consecutively. Extra bytes between entries will also be
+   removed. Data before the first referenced entry is preserved unless it
+   appears to be a sequence of consecutive local file entries.
 
    ``strict_descriptor=True`` can be provided to skip the slower scan for an
    unsigned data descriptor (deprecated in the latest ZIP specification and is

Lib/test/test_zipfile/test_core.py

@@ -2092,6 +2092,118 @@ def test_repack_overlapping_blocks(self):
                     with self.assertRaises(zipfile.BadZipFile):
                         zh.repack()
 
+    def test_repack_removed_basic(self):
+        """Should remove local file entries for provided deleted files."""
+        ln = len(self.test_files)
+        iii = (ii for n in range(1, ln + 1) for ii in itertools.combinations(range(ln), n))
+        for ii in iii:
+            with self.subTest(remove=ii):
+                # calculate the expected results
+                test_files = [data for j, data in enumerate(self.test_files) if j not in ii]
+                expected_zinfos = self._prepare_zip_from_test_files(TESTFN, test_files)
+                expected_size = os.path.getsize(TESTFN)
+
+                # do the removal and check the result
+                zinfos = self._prepare_zip_from_test_files(TESTFN, self.test_files)
+                with zipfile.ZipFile(TESTFN, 'a', self.compression) as zh:
+                    zinfos = [zh.remove(self.test_files[i][0]) for i in ii]
+                    zh.repack(zinfos)
+
+                    # check infolist
+                    self.assertEqual(
+                        [ComparableZipInfo(zi) for zi in zh.infolist()],
+                        expected_zinfos,
+                    )
+
+                # check file size
+                self.assertEqual(os.path.getsize(TESTFN), expected_size)
+
+                # make sure the zip file is still valid
+                with zipfile.ZipFile(TESTFN) as zh:
+                    self.assertIsNone(zh.testzip())
+
+    def test_repack_removed_partial(self):
+        """Should remove local file entries only for provided deleted files."""
+        ln = len(self.test_files)
+        iii = (ii for n in range(1, ln + 1) for ii in itertools.combinations(range(ln), n))
+        for ii in iii:
+            with self.subTest(removed=ii):
+                # calculate the expected results
+                test_files = [data for j, data in enumerate(self.test_files) if j not in ii]
+                expected_zinfos = self._prepare_zip_from_test_files(TESTFN, test_files)
+                with zipfile.ZipFile(TESTFN, 'a', self.compression) as zh:
+                    for zi in zh.infolist().copy():
+                        zh.remove(zi)
+                expected_size = os.path.getsize(TESTFN)
+
+                # do the removal and check the result
+                zinfos = self._prepare_zip_from_test_files(TESTFN, self.test_files)
+                with zipfile.ZipFile(TESTFN, 'a', self.compression) as zh:
+                    zinfos = [zh.remove(self.test_files[i][0]) for i, _ in enumerate(self.test_files)]
+                    zh.repack([zinfos[i] for i in ii])
+
+                    # check infolist
+                    self.assertEqual(
+                        [ComparableZipInfo(zi) for zi in zh.infolist()],
+                        [],
+                    )
+
+                # check file size
+                self.assertEqual(os.path.getsize(TESTFN), expected_size)
+
+                # make sure the zip file is still valid
+                with zipfile.ZipFile(TESTFN) as zh:
+                    self.assertIsNone(zh.testzip())
+
+    def test_repack_removed_bytes_between_files(self):
+        """Should not remove bytes between local file entries."""
+        # calculate the expected results
+        for ii in ([0], [1], [2]):
+            with self.subTest(removed=ii):
+                expected_zinfos = []
+                with open(TESTFN, 'wb') as fh:
+                    with zipfile.ZipFile(fh, 'w', self.compression) as zh:
+                        for j, (file, data) in enumerate(self.test_files):
+                            if j not in ii:
+                                zh.writestr(file, data)
+                                expected_zinfos.append(ComparableZipInfo(zh.getinfo(file)))
+                            fh.write(b' dummy bytes ')
+                expected_size = os.path.getsize(TESTFN)
+
+                # do the removal and check the result
+                with open(TESTFN, 'wb') as fh:
+                    with zipfile.ZipFile(fh, 'w', self.compression) as zh:
+                        for i, (file, data) in enumerate(self.test_files):
+                            zh.writestr(file, data)
+                            fh.write(b' dummy bytes ')
+                with zipfile.ZipFile(TESTFN, 'a', self.compression) as zh:
+                    zinfos = [zh.remove(self.test_files[i][0]) for i in ii]
+                    zh.repack(zinfos)
+
+                    # check infolist
+                    self.assertEqual(
+                        [ComparableZipInfo(zi) for zi in zh.infolist()],
+                        expected_zinfos,
+                    )
+
+                # check file size
+                self.assertEqual(os.path.getsize(TESTFN), expected_size)
+
+                # make sure the zip file is still valid
+                with zipfile.ZipFile(TESTFN) as zh:
+                    self.assertIsNone(zh.testzip())
+
+    def test_repack_removed_bad_removed_zinfos(self):
+        """Should raise when providing non-removed zinfos."""
+        # calculate the expected results
+        for ii in ([0], [1], [2]):
+            with self.subTest(removed=ii):
+                self._prepare_zip_from_test_files(TESTFN, self.test_files)
+                with zipfile.ZipFile(TESTFN, 'a') as zh:
+                    zinfos = [zh.getinfo(self.test_files[i][0]) for i in ii]
+                    with self.assertRaises(zipfile.BadZipFile):
+                        zh.repack(zinfos)
+
     @mock.patch('zipfile._ZipRepacker')
     def test_repack_closed(self, m_repack):
         self._prepare_zip_from_test_files(TESTFN, self.test_files)

Lib/zipfile/__init__.py

@@ -1378,7 +1378,7 @@ def _debug(self, level, *msg):
         if self.debug >= level:
             print(*msg)
 
-    def repack(self, zfile):
+    def repack(self, zfile, removed=None):
         """
         Repack the ZIP file, removing unrecorded local file entries and random
         bytes not listed in the central directory.
@@ -1442,11 +1442,14 @@ def repack(self, zfile):
             [recorded local file entry 1]
             ...
         """
+        removed_zinfos = set(removed or ())
+
         fp = zfile.fp
 
         # get a sorted filelist by header offset, in case the dir order
         # doesn't match the actual entry order
-        filelist = sorted(zfile.filelist, key=lambda x: x.header_offset)
+        filelist = (*zfile.filelist, *removed_zinfos)
+        filelist = sorted(filelist, key=lambda x: x.header_offset)
 
         # calculate each entry size and validate
         entry_size_list = []
@@ -1464,36 +1467,55 @@ def repack(self, zfile):
             self._debug(3, i, zinfo.orig_filename, entry_size, used_entry_size)
             if used_entry_size > entry_size:
                 raise BadZipFile(
-                    f"Overlapped entries: {zinfo.orig_filename!r} "
-                    f"(possible zip bomb)")
+                    f"Overlapped entries: {zinfo.orig_filename!r} ")
 
             entry_size_list.append(entry_size)
             used_entry_size_list.append(used_entry_size)
 
         # calculate the starting entry offset (bytes to skip)
-        try:
-            data_offset = filelist[0].header_offset
-        except IndexError:
-            data_offset = zfile.start_dir
-        entry_offset = self._calc_initial_entry_offset(fp, data_offset)
+        if removed is None:
+            try:
+                data_offset = filelist[0].header_offset
+            except IndexError:
+                data_offset = zfile.start_dir
+            entry_offset = self._calc_initial_entry_offset(fp, data_offset)
+        else:
+            entry_offset = 0
 
         # move file entries
         for i, zinfo in enumerate(filelist):
             entry_size = entry_size_list[i]
             used_entry_size = used_entry_size_list[i]
 
             # update the header and move entry data to the new position
-            if entry_offset > 0:
+            if zinfo in removed_zinfos:
                 old_header_offset = zinfo.header_offset
                 zinfo.header_offset -= entry_offset
-                self._move_entry_data(fp, old_header_offset, zinfo.header_offset, used_entry_size)
+                self._move_entry_data(
+                    fp,
+                    old_header_offset + used_entry_size,
+                    zinfo.header_offset,
+                    entry_size - used_entry_size
+                )
 
-            if zinfo._end_offset is not None:
-                zinfo._end_offset = zinfo.header_offset + used_entry_size
+                if zinfo._end_offset is not None:
+                    zinfo._end_offset = zinfo.header_offset
 
-            # update entry_offset for subsequent files to follow
-            if used_entry_size < entry_size:
-                entry_offset += entry_size - used_entry_size
+                # update entry_offset for subsequent files to follow
+                entry_offset += used_entry_size
+
+            else:
+                if entry_offset > 0:
+                    old_header_offset = zinfo.header_offset
+                    zinfo.header_offset -= entry_offset
+                    self._move_entry_data(fp, old_header_offset, zinfo.header_offset, used_entry_size)
+
+                if zinfo._end_offset is not None:
+                    zinfo._end_offset = zinfo.header_offset + used_entry_size
+
+                # update entry_offset for subsequent files to follow
+                if used_entry_size < entry_size:
+                    entry_offset += entry_size - used_entry_size
 
         # update state
         zfile.start_dir -= entry_offset
@@ -2250,7 +2272,9 @@ def remove(self, zinfo_or_arcname):
 
             self._didModify = True
 
-    def repack(self, **opts):
+        return zinfo
+
+    def repack(self, removed=None, **opts):
         """Repack a zip file, removing non-referenced file entries.
 
         The archive must be opened with mode 'a', as mode 'w'/'x' do not
@@ -2270,7 +2294,7 @@ def repack(self, **opts):
         with self._lock:
             self._writing = True
             try:
-                _ZipRepacker(**opts).repack(self)
+                _ZipRepacker(**opts).repack(self, removed)
             finally:
                 self._writing = False