Merge pull request #24 from santi921/fix/out-files-zip-overwrite

fix: preserve out_files.zip across restart / parse_only reruns

Author: santi921

qtaim_gen/source/core/omol.py

@@ -47,6 +47,7 @@
     pull_ecp_dict,
     overwrite_molden_w_ecp,
     check_spin,
+    merge_zip_into,
 )
 
 
@@ -1063,15 +1064,13 @@ def clean_jobs(
                 os.remove(os.path.join(folder, file))
                 logger.info(f"Zipped and removed {file}")
 
-        if move_results:
-            results_folder = os.path.join(folder, "generator")
-            if not os.path.exists(results_folder):
-                os.mkdir(results_folder)
-            os.rename(
-                zip_file_out,
-                os.path.join(results_folder, "out_files.zip"),
-            )
-            logger.info(f"Moved zipped out files to results folder")
+    if move_results:
+        results_folder = os.path.join(folder, "generator")
+        merge_zip_into(
+            zip_file_out,
+            os.path.join(results_folder, "out_files.zip"),
+            logger=logger,
+        )
 
 
 def setup_logger(folder: str, name: str = "gbw_analysis") -> logging.Logger:

qtaim_gen/source/core/workflow.py

@@ -416,11 +416,13 @@ def process_folder_alcf(
                                     logger.info(f"Zipped and removed {file}")
 
                     if move_results:
+                        from qtaim_gen.source.utils.io import merge_zip_into
                         results_folder = os.path.join(folder, "generator")
-                        if not os.path.exists(results_folder):
-                            os.makedirs(results_folder)
-                        shutil.move(zip_file_out, os.path.join(results_folder, "out_files.zip"))
-                        logger.info(f"Moved out_files.zip to results folder {results_folder}")
+                        merge_zip_into(
+                            zip_file_out,
+                            os.path.join(results_folder, "out_files.zip"),
+                            logger=logger,
+                        )
 
                 except Exception as e:
                     logger.info(f"Couldn't zip .out files in {folder}: {e}")

qtaim_gen/source/utils/io.py

@@ -298,6 +298,93 @@ def check_ecp_for_folder(folder_outputs: str) -> int:
         return ECP_NO_ZIP
 
 
+def merge_zip_into(
+    src_zip: str,
+    dest_zip: str,
+    logger: Optional[Any] = None,
+) -> None:
+    """Move src_zip to dest_zip, merging if dest already exists.
+
+    On filename collision, keeps whichever entry has the larger uncompressed
+    size (preserves the richest available output). On equal size, keeps the
+    existing dest entry. The merge is written to a temp file and atomically
+    replaces dest, so a failure mid-merge cannot corrupt dest. src_zip is
+    removed on success.
+    """
+    import shutil as _shutil
+
+    if not os.path.exists(src_zip):
+        if logger is not None:
+            logger.warning(f"merge_zip_into: source {src_zip} missing, nothing to do")
+        return
+
+    dest_dir = os.path.dirname(dest_zip)
+    if dest_dir:
+        os.makedirs(dest_dir, exist_ok=True)
+
+    if not os.path.exists(dest_zip):
+        _shutil.move(src_zip, dest_zip)
+        if logger is not None:
+            logger.info(f"Moved {src_zip} to {dest_zip}")
+        return
+
+    tmp_dest = dest_zip + ".merge.tmp"
+    if os.path.exists(tmp_dest):
+        try:
+            os.remove(tmp_dest)
+        except OSError:
+            pass
+
+    try:
+        with zipfile.ZipFile(dest_zip, "r") as dest_zf, \
+                zipfile.ZipFile(src_zip, "r") as src_zf:
+            dest_infos = {i.filename: i for i in dest_zf.infolist()}
+            src_infos = {i.filename: i for i in src_zf.infolist()}
+            names = set(dest_infos) | set(src_infos)
+            added: List[str] = []
+            replaced: List[Tuple[str, int, int]] = []
+            kept: List[str] = []
+            with zipfile.ZipFile(tmp_dest, "w", zipfile.ZIP_DEFLATED) as out_zf:
+                for name in sorted(names):
+                    in_dest = name in dest_infos
+                    in_src = name in src_infos
+                    if in_dest and in_src:
+                        d_sz = dest_infos[name].file_size
+                        s_sz = src_infos[name].file_size
+                        if s_sz > d_sz:
+                            out_zf.writestr(src_infos[name], src_zf.read(name))
+                            replaced.append((name, d_sz, s_sz))
+                        else:
+                            out_zf.writestr(dest_infos[name], dest_zf.read(name))
+                            kept.append(name)
+                    elif in_dest:
+                        out_zf.writestr(dest_infos[name], dest_zf.read(name))
+                        kept.append(name)
+                    else:
+                        out_zf.writestr(src_infos[name], src_zf.read(name))
+                        added.append(name)
+        os.replace(tmp_dest, dest_zip)
+        try:
+            os.remove(src_zip)
+        except OSError:
+            pass
+        if logger is not None:
+            logger.info(
+                f"Merged {src_zip} into {dest_zip}: "
+                f"{len(added)} added, {len(replaced)} replaced by larger, "
+                f"{len(kept)} kept from existing"
+            )
+            if replaced:
+                logger.info(f"Replaced (name, dest_size, src_size): {replaced}")
+    except Exception:
+        if os.path.exists(tmp_dest):
+            try:
+                os.remove(tmp_dest)
+            except OSError:
+                pass
+        raise
+
+
 def pull_ecp_dict(orca_out: str) -> Dict[int, Dict[str, Union[str, float]]]:
     """
     Method to pull the ecp dictionary from the orca output file.

tests/test_merge_zip_into.py

@@ -0,0 +1,132 @@
+"""Tests for merge_zip_into: preserve richest out_files.zip across reruns."""
+
+import os
+import tempfile
+import zipfile
+
+import pytest
+
+from qtaim_gen.source.utils.io import merge_zip_into
+
+
+def _make_zip(path: str, entries: dict) -> None:
+    with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for name, data in entries.items():
+            zf.writestr(name, data)
+
+
+def _read_zip(path: str) -> dict:
+    out = {}
+    with zipfile.ZipFile(path, "r") as zf:
+        for name in zf.namelist():
+            out[name] = zf.read(name)
+    return out
+
+
+def test_move_when_dest_missing():
+    with tempfile.TemporaryDirectory() as tmp:
+        src = os.path.join(tmp, "src.zip")
+        dest = os.path.join(tmp, "sub", "dest.zip")
+        _make_zip(src, {"a.out": b"hello"})
+
+        merge_zip_into(src, dest)
+
+        assert not os.path.exists(src)
+        assert os.path.exists(dest)
+        assert _read_zip(dest) == {"a.out": b"hello"}
+
+
+def test_merge_adds_missing_entries():
+    with tempfile.TemporaryDirectory() as tmp:
+        src = os.path.join(tmp, "src.zip")
+        dest = os.path.join(tmp, "dest.zip")
+        _make_zip(dest, {"a.out": b"AAA"})
+        _make_zip(src, {"b.out": b"BBB", "c.out": b"CCC"})
+
+        merge_zip_into(src, dest)
+
+        assert not os.path.exists(src)
+        assert _read_zip(dest) == {
+            "a.out": b"AAA",
+            "b.out": b"BBB",
+            "c.out": b"CCC",
+        }
+
+
+def test_collision_prefers_larger():
+    with tempfile.TemporaryDirectory() as tmp:
+        src = os.path.join(tmp, "src.zip")
+        dest = os.path.join(tmp, "dest.zip")
+        _make_zip(dest, {"adch.out": b"short", "cm5.out": b"rich cm5 content"})
+        _make_zip(src, {"adch.out": b"richer adch content here", "cm5.out": b"x"})
+
+        merge_zip_into(src, dest)
+
+        merged = _read_zip(dest)
+        assert merged["adch.out"] == b"richer adch content here"
+        assert merged["cm5.out"] == b"rich cm5 content"
+
+
+def test_collision_equal_size_keeps_existing():
+    with tempfile.TemporaryDirectory() as tmp:
+        src = os.path.join(tmp, "src.zip")
+        dest = os.path.join(tmp, "dest.zip")
+        _make_zip(dest, {"a.out": b"XXXXX"})
+        _make_zip(src, {"a.out": b"YYYYY"})
+
+        merge_zip_into(src, dest)
+
+        assert _read_zip(dest) == {"a.out": b"XXXXX"}
+
+
+def test_missing_src_is_noop():
+    with tempfile.TemporaryDirectory() as tmp:
+        src = os.path.join(tmp, "missing.zip")
+        dest = os.path.join(tmp, "dest.zip")
+        _make_zip(dest, {"a.out": b"AAA"})
+
+        merge_zip_into(src, dest)
+
+        assert _read_zip(dest) == {"a.out": b"AAA"}
+
+
+def test_dest_corrupt_raises_and_preserves_dest():
+    with tempfile.TemporaryDirectory() as tmp:
+        src = os.path.join(tmp, "src.zip")
+        dest = os.path.join(tmp, "dest.zip")
+        _make_zip(src, {"a.out": b"AAA"})
+        with open(dest, "wb") as f:
+            f.write(b"not a zip file")
+
+        with pytest.raises(zipfile.BadZipFile):
+            merge_zip_into(src, dest)
+
+        assert os.path.exists(src)
+        with open(dest, "rb") as f:
+            assert f.read() == b"not a zip file"
+        assert not os.path.exists(dest + ".merge.tmp")
+
+
+def test_simulated_rerun_preserves_original_entries():
+    # Simulate first full run producing many .out files, then a parse_only
+    # rerun producing only a subset - merged zip should retain everything.
+    with tempfile.TemporaryDirectory() as tmp:
+        dest = os.path.join(tmp, "out_files.zip")
+        first_run = {
+            "adch.out": b"first adch full output",
+            "cm5.out": b"first cm5 full output",
+            "hirshfeld.out": b"first hirshfeld full output",
+            "fuzzy_full.out": b"first fuzzy full output",
+        }
+        _make_zip(dest, first_run)
+
+        src = os.path.join(tmp, "out_files.zip.new")
+        rerun = {"adch.out": b"short rerun"}
+        _make_zip(src, rerun)
+
+        merge_zip_into(src, dest)
+
+        merged = _read_zip(dest)
+        assert set(merged) == set(first_run)
+        for k, v in first_run.items():
+            assert merged[k] == v