diff --git a/autofit/non_linear/paths/abstract.py b/autofit/non_linear/paths/abstract.py index a80cbc23d..7764724c9 100644 --- a/autofit/non_linear/paths/abstract.py +++ b/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): """ diff --git a/test_autofit/non_linear/paths/test_paths.py b/test_autofit/non_linear/paths/test_paths.py index 1837ae356..a7777da8c 100644 --- a/test_autofit/non_linear/paths/test_paths.py +++ b/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))