Update data_offset

Author: danny0838

src/zipremove/__init__.py

@@ -177,7 +177,8 @@ def repack(self, zfile, removed=None):
 
         Side effects:
             - Modifies the ZIP file in place.
-            - Updates zfile.start_dir to account for removed data.
+            - Updates zfile.start_dir and zfile.data_offset to account for
+              removed data.
             - Sets zfile._didModify to True.
             - Updates header_offset and clears _end_offset of referenced
               ZipInfo instances.
@@ -282,6 +283,15 @@ def repack(self, zfile, removed=None):
         zfile.start_dir -= entry_offset
         zfile._didModify = True
 
+        # polyfill: update _data_offset if exists
+        if getattr(zfile, '_data_offset', None):
+            try:
+                offset = filelist[0].header_offset
+            except IndexError:
+                offset = zfile.start_dir
+            if offset < zfile._data_offset:
+                zfile._data_offset = offset
+
         # polyfill: clear ZipInfo._end_offset if exists
         # (Python >= 3.8 with fix #109858)
         if hasattr(ZipInfo, '_end_offset'):

tests/test_zipfile.py

@@ -966,8 +966,8 @@ def test_repack_prepended_bytes(self):
         for ii in ([], [0], [0, 1], [1], [2]):
             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]
                 fz = io.BytesIO()
+                test_files = [data for j, data in enumerate(self.test_files) if j not in ii]
                 self._prepare_zip_from_test_files(fz, test_files)
                 fz.seek(0)
                 with open(TESTFN, 'wb') as fh:
@@ -985,9 +985,67 @@ def test_repack_prepended_bytes(self):
                     fh.write(b'dummy ')
                     fh.write(fz.read())
                 with zipfile.ZipFile(TESTFN, 'a', self.compression) as zh:
+                    if hasattr(zh, 'data_offset'):
+                        self.assertEqual(zh.data_offset, 6)
                     for i in ii:
                         zh.remove(self.test_files[i][0])
                     zh.repack()
+                    if hasattr(zh, 'data_offset'):
+                        self.assertEqual(zh.data_offset, 6)
+
+                    # check infolist
+                    self.assertEqual(
+                        [ComparableZipInfo(zi) for zi in zh.infolist()],
+                        [ComparableZipInfo(zi) for zi in 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_prepended_file_entry(self):
+        for ii in ([0], [0, 1], [0, 1, 2]):
+            with self.subTest(remove=ii):
+                # calculate the expected results
+                fz = io.BytesIO()
+                test_files = [data for j, data in enumerate(self.test_files) if j not in ii]
+                self._prepare_zip_from_test_files(fz, test_files)
+                fz.seek(0)
+                with open(TESTFN, 'wb') as fh:
+                    fh.write(b'dummy ')
+                    fh.write(fz.read())
+                with zipfile.ZipFile(TESTFN) as zh:
+                    expected_zinfos = list(zh.infolist())
+                expected_size = os.path.getsize(TESTFN)
+
+                # do the removal and check the result
+                fz = io.BytesIO()
+                with zipfile.ZipFile(fz, 'w') as zh:
+                    for j, (file, data) in enumerate(self.test_files):
+                        if j in ii:
+                            zh.writestr(file, data)
+                    fz.seek(0)
+                    prefix = fz.read()
+
+                fz = io.BytesIO()
+                test_files = [data for j, data in enumerate(self.test_files) if j not in ii]
+                self._prepare_zip_from_test_files(fz, test_files)
+                fz.seek(0)
+
+                with open(TESTFN, 'wb') as fh:
+                    fh.write(b'dummy ')
+                    fh.write(prefix)
+                    fh.write(fz.read())
+
+                with zipfile.ZipFile(TESTFN, 'a', self.compression) as zh:
+                    if hasattr(zh, 'data_offset'):
+                        self.assertEqual(zh.data_offset, 6 + len(prefix))
+                    zh.repack()
+                    if hasattr(zh, 'data_offset'):
+                        self.assertEqual(zh.data_offset, 6)
 
                     # check infolist
                     self.assertEqual(
@@ -1177,8 +1235,12 @@ def test_repack_removed_prepended_bytes(self):
                     fh.write(b'dummy ')
                     fh.write(fz.read())
                 with zipfile.ZipFile(TESTFN, 'a', self.compression) as zh:
+                    if hasattr(zh, 'data_offset'):
+                        self.assertEqual(zh.data_offset, 6)
                     zinfos = [zh.remove(self.test_files[i][0]) for i in ii]
                     zh.repack(zinfos)
+                    if hasattr(zh, 'data_offset'):
+                        self.assertEqual(zh.data_offset, 6)
 
                     # check infolist
                     self.assertEqual(