Prevent util.execute from reading output files

qcengine/tests/test_utils.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import time
+import pathlib
 
 import pytest
 from qcelemental.models import AtomicInput
@@ -64,17 +65,29 @@ def test_tmpdir():
         assert str(tmpdir).split(os.path.sep)[-1].endswith("this")
 
 
-def test_disk_files():
+@pytest.mark.parametrize("outfiles_track", [[], ["thing*"], ["thing*", "other"]])
+def test_disk_files(outfiles_track):
 
     infiles = {"thing1": "hello", "thing2": "world", "other": "everyone"}
     outfiles = {"thing*": None, "other": None}
     with util.temporary_directory(suffix="this") as tmpdir:
-        with util.disk_files(infiles=infiles, outfiles=outfiles, cwd=tmpdir):
+        with util.disk_files(infiles=infiles, outfiles=outfiles, cwd=tmpdir, outfiles_track=outfiles_track):
             pass
 
     assert outfiles.keys() == {"thing*", "other"}
-    assert outfiles["thing*"]["thing1"] == "hello"
-    assert outfiles["other"] == "everyone"
+    for ofile, ofile_val in outfiles.items():
+        if isinstance(ofile_val, dict):
+            if ofile in outfiles_track:
+                for fpath in ofile_val.values():
+                    assert isinstance(fpath, pathlib.PurePath)
+            else:
+                for key in ofile_val.keys():
+                    print(key)
+                    assert ofile_val[key] == infiles[key]
+        elif ofile in outfiles_track:
+            assert isinstance(ofile_val, pathlib.PurePath)
+        else:
+            assert ofile_val == infiles[ofile]
 
 
 def test_popen_tee_output(capsys):

qcengine/util.py

@@ -377,6 +377,7 @@ def execute(
     infiles: Optional[Dict[str, str]] = None,
     outfiles: Optional[List[str]] = None,
     *,
+    outfiles_track: Optional[List[str]] = [],
     as_binary: Optional[List[str]] = None,
     scratch_name: Optional[str] = None,
     scratch_directory: Optional[str] = None,
@@ -404,6 +405,11 @@ def execute(
     outfiles : List[str] = None
         Output file names to be collected after execution into
         values. May be {}.
+    outfiles_track: List[str], optional
+        Keys of `outfiles` to keep track of without loading their contents in memory.
+        For specified filename in `outfiles_track`, the file path instead of contents
+        is stored in `outfiles`. To ensure tracked files are not deleted after execution,
+        you must set `scratch_messy=True`.
     as_binary : List[str] = None
         Keys of `infiles` or `outfiles` to be treated as bytes.
     scratch_name : str, optional
@@ -428,7 +434,6 @@ def execute(
         Run command through the shell.
     exit_code: int, optional
         The exit code above which the process is considered failure.
-
     Raises
     ------
     FileExistsError
@@ -471,7 +476,9 @@ def execute(
     ) as scrdir:
         popen_kwargs["cwd"] = scrdir
         popen_kwargs["shell"] = shell
-        with disk_files(infiles, outfiles, cwd=scrdir, as_binary=as_binary) as extrafiles:
+        with disk_files(
+            infiles, outfiles, cwd=scrdir, as_binary=as_binary, outfiles_track=outfiles_track
+        ) as extrafiles:
             with popen(command, popen_kwargs=popen_kwargs) as proc:
                 # Wait for the subprocess to complete or the timeout to expire
                 if interupt_after is None:
@@ -562,7 +569,8 @@ def disk_files(
     *,
     cwd: Optional[str] = None,
     as_binary: Optional[List[str]] = None,
-) -> Dict[str, Union[str, bytes]]:
+    outfiles_track: Optional[List[str]] = [],
+) -> Dict[str, Union[str, bytes, Path]]:
     """Write and collect files.
 
     Parameters
@@ -577,10 +585,13 @@ def disk_files(
         Directory to which to write and read files.
     as_binary : List[str] = None
         Keys in `infiles` (`outfiles`) to be written (read) as bytes, not decoded.
-
+    outfiles_track: List[str], optional
+        Keys of `outfiles` to keep track of (i.e. file contents not loaded in memory).
+        For specified filename in `outfiles_track`, the file path instead of contents
+        is stored in `outfiles`.
     Yields
     ------
-    Dict[str] = str
+    Dict[str, Union[str, bytes, Path]]
         outfiles with RHS filled in.
 
     """
@@ -604,19 +615,33 @@ def disk_files(
 
     finally:
         for fl in outfiles.keys():
-            omode = "rb" if fl in as_binary else "r"
-            try:
-                filename = lwd / fl
-                with open(filename, omode) as fp:
-                    outfiles[fl] = fp.read()
-                    LOGGER.info(f"... Writing ({omode}): {filename}")
-            except (OSError, FileNotFoundError):
-                if "*" in fl:
+            filename = lwd / fl
+            if fl not in outfiles_track:
+                omode = "rb" if fl in as_binary else "r"
+                try:
+                    with open(filename, omode) as fp:
+                        outfiles[fl] = fp.read()
+                        LOGGER.info(f"... Reading ({omode}): {filename}")
+                except (OSError, FileNotFoundError):
+                    if "*" in fl:
+                        gfls = {}
+                        for gfl in lwd.glob(fl):
+                            with open(gfl, omode) as fp:
+                                gfls[gfl.name] = fp.read()
+                                LOGGER.info(f"... Reading ({omode}): {gfl}")
+                        if not gfls:
+                            gfls = None
+                        outfiles[fl] = gfls
+                    else:
+                        outfiles[fl] = None
+            else:
+                if filename.is_file():
+                    outfiles[fl] = filename
+                elif "*" in fl:
                     gfls = {}
                     for gfl in lwd.glob(fl):
-                        with open(gfl, omode) as fp:
-                            gfls[gfl.name] = fp.read()
-                            LOGGER.info(f"... Writing ({omode}): {gfl}")
+                        if gfl.is_file():
+                            gfls[gfl.name] = gfl
                     if not gfls:
                         gfls = None
                     outfiles[fl] = gfls