fix: preserve_in_zip replaces a stale archive member (#1414)

autofit/non_linear/paths/abstract.py

@@ -3,6 +3,7 @@
 import re
 import shutil
 import zipfile
+import zlib
 from abc import ABC, abstractmethod
 from configparser import NoSectionError
 from pathlib import Path
@@ -36,6 +37,72 @@ def _test_mode_segment() -> Optional[str]:
     return "test_mode" if is_test_mode() else None
 
 
+def _matches_archived(info: zipfile.ZipInfo, file_path) -> bool:
+    """
+    Returns whether the archived member described by ``info`` has the same
+    content as the file at ``file_path``.
+
+    Compared via the size and CRC already held in the zip's central directory,
+    so the archived bytes are never decompressed and the file on disk is read
+    once in chunks — a cache artifact can be a large ``.fits`` image.
+    """
+    if info.file_size != os.path.getsize(file_path):
+        return False
+
+    crc = 0
+    with open(file_path, "rb") as f:
+        for chunk in iter(lambda: f.read(1024 * 1024), b""):
+            crc = zlib.crc32(chunk, crc)
+
+    return crc == info.CRC
+
+
+def _replace_zip_member(zip_path, arcname: str, file_path):
+    """
+    Replace the member ``arcname`` of the zip at ``zip_path`` with the file at
+    ``file_path``.
+
+    ``zipfile`` cannot overwrite a member in place, so every other member is
+    streamed into a new archive alongside the replacement. The new archive is
+    written next to the original (same directory, so the final ``os.replace``
+    is atomic on one filesystem) and only swapped in once complete: an error
+    part-way leaves the original archive untouched.
+
+    Rebuilding the archive from the output directory instead (as the search
+    itself does at completion) is not an option here, because with
+    ``remove_files`` the directory holds only the file just written.
+    """
+    zip_path = Path(zip_path)
+    temporary_path = zip_path.with_name(f"{zip_path.name}.replace.tmp")
+
+    try:
+        with zipfile.ZipFile(zip_path, "r") as source:
+            replaced = source.getinfo(arcname)
+
+            with zipfile.ZipFile(temporary_path, "w") as destination:
+                for info in source.infolist():
+                    if info.filename == arcname:
+                        continue
+
+                    with source.open(info) as member:
+                        with destination.open(info, "w") as target:
+                            shutil.copyfileobj(member, target)
+
+                destination.write(
+                    file_path,
+                    arcname,
+                    compress_type=replaced.compress_type,
+                )
+
+        os.replace(temporary_path, zip_path)
+    except BaseException:
+        try:
+            os.remove(temporary_path)
+        except FileNotFoundError:
+            pass
+        raise
+
+
 class AbstractPaths(ABC):
     def __init__(
         self,
@@ -325,7 +392,10 @@ def preserve_in_zip(self, file_path):
 
         No-op when the zip does not exist (e.g. the search is still running,
         so the file will be zipped with everything else at completion) and
-        when the member is already present.
+        when the archived member is already byte-identical to the file on
+        disk. When the member exists but its content has changed — a cache
+        that was invalidated and recomputed — it is replaced, otherwise the
+        stale copy would come back at the next ``restore()``.
 
         Parameters
         ----------
@@ -339,8 +409,20 @@ def preserve_in_zip(self, file_path):
         arcname = str(Path(file_path).relative_to(self.output_path))
 
         with zipfile.ZipFile(self._zip_path, "a") as f:
-            if arcname not in f.namelist():
+            try:
+                info = f.getinfo(arcname)
+            except KeyError:
                 f.write(file_path, arcname)
+                return
+
+            if _matches_archived(info, file_path):
+                return
+
+        _replace_zip_member(
+            zip_path=self._zip_path,
+            arcname=arcname,
+            file_path=file_path,
+        )
 
     def restore(self):
         """

test_autofit/non_linear/paths/test_paths.py

@@ -149,6 +149,59 @@ def test__preserve_in_zip__file_survives_restore(tmp_path):
     assert cache_file.exists()
 
 
+def test__preserve_in_zip__replaces_stale_member(tmp_path):
+    import zipfile
+
+    paths = af.DirectoryPaths(name="preserve_replace_test", path_prefix=str(tmp_path))
+
+    files_path = Path(paths._files_path)
+    files_path.mkdir(parents=True, exist_ok=True)
+    (files_path / "samples_summary.json").write_text("{}")
+
+    cache_file = files_path / "cache_artifact.json"
+    cache_file.write_text('{"cached": "stale"}')
+
+    paths.zip_remove()
+    assert Path(paths._zip_path).exists()
+
+    # The cache was invalidated and recomputed: the member is already in the
+    # zip, but its content has changed.
+    files_path.mkdir(parents=True, exist_ok=True)
+    cache_file.write_text('{"cached": "recomputed"}')
+    paths.preserve_in_zip(cache_file)
+
+    with zipfile.ZipFile(paths._zip_path) as f:
+        assert f.namelist().count("files/cache_artifact.json") == 1
+        assert f.read("files/cache_artifact.json") == b'{"cached": "recomputed"}'
+        # The rewrite must not lose the members it copied across.
+        assert f.read("files/samples_summary.json") == b"{}"
+
+    # The restore cycle yields the recomputed bytes, not the stale ones.
+    paths.restore()
+    assert cache_file.read_text() == '{"cached": "recomputed"}'
+
+
+def test__preserve_in_zip__identical_content_does_not_rewrite(tmp_path):
+    paths = af.DirectoryPaths(name="preserve_identical_test", path_prefix=str(tmp_path))
+
+    files_path = Path(paths._files_path)
+    files_path.mkdir(parents=True, exist_ok=True)
+    cache_file = files_path / "cache_artifact.json"
+    cache_file.write_text('{"cached": true}')
+
+    paths.zip_remove()
+    before = Path(paths._zip_path).read_bytes()
+
+    # Re-preserving a file whose content is unchanged leaves the archive
+    # untouched, so the common resume path never pays for a rewrite.
+    files_path.mkdir(parents=True, exist_ok=True)
+    cache_file.write_text('{"cached": true}')
+
+    paths.preserve_in_zip(cache_file)
+
+    assert Path(paths._zip_path).read_bytes() == before
+
+
 def test__preserve_in_zip__no_zip_is_a_no_op(tmp_path):
     paths = af.DirectoryPaths(name="preserve_noop_test", path_prefix=str(tmp_path))