diff --git a/autofit/non_linear/paths/abstract.py b/autofit/non_linear/paths/abstract.py index d1d40426d..0485535f7 100644 --- a/autofit/non_linear/paths/abstract.py +++ b/autofit/non_linear/paths/abstract.py @@ -313,6 +313,35 @@ def _zip(self): except FileNotFoundError: pass + def preserve_in_zip(self, file_path): + """ + Add a file (already inside this search's output directory) to the + search's ``.zip`` archive so it survives the resume cycle. + + ``restore()`` deletes the output directory and re-extracts the zip, so + any file written into e.g. ``files/`` *after* a search completed — such + as a cache artifact derived from the finished result — would be + destroyed by the next resume unless it is also a member of the zip. + + 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. + + Parameters + ---------- + file_path + Absolute path of the file to preserve; must live under + ``output_path``, from which its archive name is derived. + """ + if not Path(self._zip_path).exists(): + return + + 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(): + f.write(file_path, arcname) + def restore(self): """ Copy files from the ``.zip`` file to the samples folder. diff --git a/test_autofit/non_linear/paths/test_paths.py b/test_autofit/non_linear/paths/test_paths.py index dd875caa4..50635bf20 100644 --- a/test_autofit/non_linear/paths/test_paths.py +++ b/test_autofit/non_linear/paths/test_paths.py @@ -115,3 +115,49 @@ def test_make_path_excludes_segment_when_level_zero(self, monkeypatch): monkeypatch.setenv("PYAUTO_TEST_MODE", "0") paths = af.DirectoryPaths(name="name", path_prefix="prefix") assert "test_mode" not in paths._make_path().parts + + +def test__preserve_in_zip__file_survives_restore(tmp_path): + import zipfile + + paths = af.DirectoryPaths(name="preserve_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("{}") + + paths.zip_remove() + assert Path(paths._zip_path).exists() + + # A post-completion cache write: on disk but not in the zip. + cache_file = files_path / "cache_artifact.json" + files_path.mkdir(parents=True, exist_ok=True) + cache_file.write_text('{"cached": true}') + + paths.preserve_in_zip(cache_file) + + with zipfile.ZipFile(paths._zip_path) as f: + assert "files/cache_artifact.json" in f.namelist() + + # Idempotent — appending again must not duplicate the member. + paths.preserve_in_zip(cache_file) + with zipfile.ZipFile(paths._zip_path) as f: + assert f.namelist().count("files/cache_artifact.json") == 1 + + # The restore cycle (rmtree + re-extract) keeps the preserved file. + paths.restore() + assert cache_file.exists() + + +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)) + + 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.preserve_in_zip(cache_file) + + assert not Path(paths._zip_path).exists() + assert cache_file.exists()