Merge pull request #203 from djarecka/mnt/wf_copyfile

[mnt] copy files to the wf directory (closes # 145)

Author: djarecka

pydra/engine/helpers.py

@@ -10,7 +10,7 @@
 import subprocess as sp
 
 from .specs import Runtime, File, attr_fields
-from .helpers_file import is_existing_file, hash_file
+from .helpers_file import is_existing_file, hash_file, copyfile, is_existing_file
 
 
 def ensure_list(obj, tuple2list=False):
@@ -112,13 +112,42 @@ def save(task_path: Path, result=None, task=None):
         raise ValueError("Nothing to be saved")
     task_path.mkdir(parents=True, exist_ok=True)
     if result:
+        if Path(task_path).name.startswith("Workflow"):
+            # copy files to the workflow directory
+            result = copyfile_workflow(wf_path=task_path, result=result)
         with (task_path / "_result.pklz").open("wb") as fp:
             cp.dump(result, fp)
     if task:
         with (task_path / "_task.pklz").open("wb") as fp:
             cp.dump(task, fp)
 
 
+def copyfile_workflow(wf_path, result):
+    """ if file in the wf results, the file will be copied to the workflow directory"""
+    for field in attr_fields(result.output):
+        value = getattr(result.output, field.name)
+        new_value = _copyfile_single_value(wf_path=wf_path, value=value)
+        if new_value != value:
+            setattr(result.output, field.name, new_value)
+    return result
+
+
+def _copyfile_single_value(wf_path, value):
+    """ checking a single value for files that need to be copied to the wf dir"""
+    if isinstance(value, (tuple, list)):
+        return [_copyfile_single_value(wf_path, val) for val in value]
+    elif isinstance(value, dict):
+        return {
+            key: _copyfile_single_value(wf_path, val) for (key, val) in value.items()
+        }
+    elif is_existing_file(value):
+        new_path = wf_path / Path(value).name
+        copyfile(originalfile=value, newfile=new_path, copy=True, use_hardlink=True)
+        return new_path
+    else:
+        return value
+
+
 def task_hash(task):
     """
     Calculate the checksum of a task.

pydra/engine/helpers_file.py

@@ -516,6 +516,8 @@ def is_local_file(f):
 
 def is_existing_file(f):
     """ checking if an object is an existing file"""
+    if not f:
+        return False
     try:
         return Path(f).exists()
     except TypeError:

pydra/engine/tests/test_shelltask.py

@@ -2018,6 +2018,8 @@ def test_shell_cmd_outputspec_wf_1(plugin):
     res = wf.result()
     assert res.output.stdout == ""
     assert res.output.newfile.exists()
+    # checking if the file was copied to the wf dir
+    assert res.output.newfile.parent == wf.output_dir
 
 
 def no_fsl():

pydra/engine/tests/test_workflow.py

@@ -14,6 +14,9 @@
     add2_sub2_res,
     fun_addvar_none,
     fun_addvar_default,
+    fun_write_file,
+    fun_write_file_list,
+    fun_write_file_list2dict,
 )
 from ..submitter import Submitter
 from ..core import Workflow
@@ -2825,3 +2828,72 @@ def test_wf_lzoutall_st_2a(plugin):
         {"out_add": [8, 62], "out_sub": [4, 58]},
         {"out_add": [62, 602], "out_sub": [58, 598]},
     ]
+
+
+# worfklows that have files in the result, the files should be copied to the wf dir
+
+
+@pytest.mark.parametrize("plugin", Plugins)
+def test_wf_resultfile_1(plugin):
+    """ workflow with a file in the result, file should be copied to the wf dir"""
+    wf = Workflow(name="wf_file_1", input_spec=["x"])
+    wf.add(fun_write_file(name="writefile", filename=wf.lzin.x))
+    wf.inputs.x = "file_1.txt"
+    wf.plugin = plugin
+    wf.set_output([("wf_out", wf.writefile.lzout.out)])
+
+    with Submitter(plugin=plugin) as sub:
+        sub(wf)
+
+    results = wf.result()
+    # checking if the file exists and if it is in the Workflow directory
+    assert results.output.wf_out.exists()
+    assert results.output.wf_out == wf.output_dir / "file_1.txt"
+
+
+@pytest.mark.parametrize("plugin", Plugins)
+def test_wf_resultfile_2(plugin):
+    """ workflow with a list of files in the wf result,
+        all files should be copied to the wf dir
+    """
+    wf = Workflow(name="wf_file_1", input_spec=["x"])
+    wf.add(fun_write_file_list(name="writefile", filename_list=wf.lzin.x))
+    file_list = ["file_1.txt", "file_2.txt", "file_3.txt"]
+    wf.inputs.x = file_list
+    wf.plugin = plugin
+    wf.set_output([("wf_out", wf.writefile.lzout.out)])
+
+    with Submitter(plugin=plugin) as sub:
+        sub(wf)
+
+    results = wf.result()
+    # checking if the file exists and if it is in the Workflow directory
+    for ii, file in enumerate(results.output.wf_out):
+        assert file.exists()
+        assert file == wf.output_dir / file_list[ii]
+
+
+@pytest.mark.parametrize("plugin", Plugins)
+def test_wf_resultfile_3(plugin):
+    """ workflow with a dictionaries of files in the wf result,
+        all files should be copied to the wf dir
+    """
+    wf = Workflow(name="wf_file_1", input_spec=["x"])
+    wf.add(fun_write_file_list2dict(name="writefile", filename_list=wf.lzin.x))
+    file_list = ["file_1.txt", "file_2.txt", "file_3.txt"]
+    wf.inputs.x = file_list
+    wf.plugin = plugin
+    wf.set_output([("wf_out", wf.writefile.lzout.out)])
+
+    with Submitter(plugin=plugin) as sub:
+        sub(wf)
+
+    results = wf.result()
+    # checking if the file exists and if it is in the Workflow directory
+    for key, val in results.output.wf_out.items():
+        if key == "random_int":
+            assert val == 20
+        else:
+            assert val.exists()
+            ii = int(key.split("_")[1])
+            assert val == wf.output_dir / file_list[ii]

pydra/engine/tests/utils.py

@@ -1,5 +1,7 @@
 # Tasks for testing
 import time
+import typing as tp
+from pathlib import Path
 
 from ..core import Workflow
 from ... import mark
@@ -102,6 +104,36 @@ def fun_dict(d):
     return "_".join(kv_list)
 
 
+@mark.task
+def fun_write_file(filename: tp.Union[str, File, Path], text="hello"):
+    with open(filename, "w") as f:
+        f.write(text)
+    return Path(filename).absolute()
+
+
+@mark.task
+def fun_write_file_list(filename_list: tp.List[tp.Union[str, File, Path]], text="hi"):
+    for ii, filename in enumerate(filename_list):
+        with open(filename, "w") as f:
+            f.write(f"from file {ii}: {text}")
+    filename_list = [Path(filename).absolute() for filename in filename_list]
+    return filename_list
+
+
+@mark.task
+def fun_write_file_list2dict(
+    filename_list: tp.List[tp.Union[str, File, Path]], text="hi"
+):
+    filename_dict = {}
+    for ii, filename in enumerate(filename_list):
+        with open(filename, "w") as f:
+            f.write(f"from file {ii}: {text}")
+        filename_dict[f"file_{ii}"] = Path(filename).absolute()
+    # adding an additional field with int
+    filename_dict["random_int"] = 20
+    return filename_dict
+
+
 @mark.task
 def fun_file(filename: File):
     with open(filename) as f: