edit to copyfile_workflow, that files from list or dictionaries are also copied

Author: djarecka

pydra/engine/helpers.py

@@ -126,13 +126,28 @@ 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)
-        if is_existing_file(value):
-            new_path = wf_path / Path(value).name
-            copyfile(originalfile=value, newfile=new_path, copy=True, use_hardlink=True)
-            setattr(result.output, field.name, new_path)
+        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/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: