Merge pull request #553 from djarecka/enh/callable_template

function in output_file_template

Author: djarecka

pydra/engine/helpers_file.py

@@ -678,7 +678,11 @@ def _template_formatting(field, inputs, inputs_dict_st):
     """
     from .specs import MultiOutputFile
 
+    # if a template is a function it has to be run first with the inputs as the only arg
     template = field.metadata["output_file_template"]
+    if callable(template):
+        template = template(inputs)
+
     # as default, we assume that keep_extension is True
     keep_extension = field.metadata.get("keep_extension", True)
 

pydra/engine/specs.py

@@ -624,7 +624,11 @@ def _check_requires(self, fld, inputs):
             # if the output has output_file_template field,
             # adding all input fields from the template to requires
             if "output_file_template" in fld.metadata:
-                inp_fields = re.findall(r"{\w+}", fld.metadata["output_file_template"])
+                template = fld.metadata["output_file_template"]
+                # if a template is a function it has to be run first with the inputs as the only arg
+                if callable(template):
+                    template = template(inputs)
+                inp_fields = re.findall(r"{\w+}", template)
                 field_required += [
                     el[1:-1] for el in inp_fields if el[1:-1] not in field_required
                 ]

pydra/engine/tests/test_shelltask.py

@@ -1556,6 +1556,72 @@ def test_shell_cmd_inputsspec_11():
         assert out_file.name == "test1" or out_file.name == "test2"
 
 
+@pytest.mark.parametrize("results_function", [result_no_submitter, result_submitter])
+def test_shell_cmd_inputspec_12(tmpdir, plugin, results_function):
+    """
+    providing output name using input_spec
+    output_file_template is provided as a function that returns
+    various templates depending on the values of inputs fields
+    """
+    cmd = "cp"
+    file = tmpdir.mkdir("data_inp").join("file.txt")
+    file.write("content\n")
+
+    def template_function(inputs):
+        if inputs.number % 2 == 0:
+            return "{file_orig}_even"
+        else:
+            return "{file_orig}_odd"
+
+    my_input_spec = SpecInfo(
+        name="Input",
+        fields=[
+            (
+                "file_orig",
+                attr.ib(
+                    type=File,
+                    metadata={"position": 2, "help_string": "new file", "argstr": ""},
+                ),
+            ),
+            (
+                "number",
+                attr.ib(
+                    type=int,
+                    metadata={"help_string": "a number", "mandatory": True},
+                ),
+            ),
+            (
+                "file_copy",
+                attr.ib(
+                    type=str,
+                    metadata={
+                        "output_file_template": template_function,
+                        "help_string": "output file",
+                        "argstr": "",
+                    },
+                ),
+            ),
+        ],
+        bases=(ShellSpec,),
+    )
+
+    shelly = ShellCommandTask(
+        name="shelly",
+        executable=cmd,
+        input_spec=my_input_spec,
+        file_orig=file,
+        number=2,
+        cache_dir=tmpdir,
+    )
+
+    res = results_function(shelly, plugin)
+    assert res.output.stdout == ""
+    assert res.output.file_copy.exists()
+    assert res.output.file_copy.name == "file_even.txt"
+    # checking if it's created in a good place
+    assert shelly.output_dir == res.output.file_copy.parent
+
+
 @pytest.mark.parametrize("results_function", [result_no_submitter, result_submitter])
 def test_shell_cmd_inputspec_copyfile_1(plugin, results_function, tmpdir):
     """shelltask changes a file in place,

pydra/engine/tests/test_shelltask_inputspec.py

@@ -1721,6 +1721,116 @@ def test_shell_cmd_inputs_template_10():
     assert shelly.output_names == ["return_code", "stdout", "stderr", "outA"]
 
 
+def test_shell_cmd_inputs_template_function_1():
+    """one input field uses output_file_template that is a simple function
+    this can be easily done by simple template as in test_shell_cmd_inputs_template_1
+    """
+
+    # a function that return an output template
+    def template_fun(inputs):
+        return "{inpA}_out"
+
+    my_input_spec = SpecInfo(
+        name="Input",
+        fields=[
+            (
+                "inpA",
+                attr.ib(
+                    type=str,
+                    metadata={
+                        "position": 1,
+                        "help_string": "inpA",
+                        "argstr": "",
+                        "mandatory": True,
+                    },
+                ),
+            ),
+            (
+                "outA",
+                attr.ib(
+                    type=str,
+                    metadata={
+                        "position": 2,
+                        "help_string": "outA",
+                        "argstr": "-o",
+                        "output_file_template": template_fun,
+                    },
+                ),
+            ),
+        ],
+        bases=(ShellSpec,),
+    )
+
+    shelly = ShellCommandTask(
+        executable="executable", input_spec=my_input_spec, inpA="inpA"
+    )
+
+    assert shelly.cmdline == f"executable inpA -o {str(shelly.output_dir / 'inpA_out')}"
+
+
+def test_shell_cmd_inputs_template_function_2():
+    """one input field uses output_file_template that is a function,
+    depending on a value of an input it returns different template
+    """
+
+    # a function that return an output template that depends on value of the input
+    def template_fun(inputs):
+        if inputs.inpB % 2 == 0:
+            return "{inpA}_even"
+        else:
+            return "{inpA}_odd"
+
+    my_input_spec = SpecInfo(
+        name="Input",
+        fields=[
+            (
+                "inpA",
+                attr.ib(
+                    type=str,
+                    metadata={
+                        "position": 1,
+                        "help_string": "inpA",
+                        "argstr": "",
+                        "mandatory": True,
+                    },
+                ),
+            ),
+            (
+                "inpB",
+                attr.ib(
+                    type=int,
+                    metadata={
+                        "help_string": "inpB",
+                        "mandatory": True,
+                    },
+                ),
+            ),
+            (
+                "outA",
+                attr.ib(
+                    type=str,
+                    metadata={
+                        "position": 2,
+                        "help_string": "outA",
+                        "argstr": "-o",
+                        "output_file_template": template_fun,
+                    },
+                ),
+            ),
+        ],
+        bases=(ShellSpec,),
+    )
+
+    shelly = ShellCommandTask(
+        executable="executable",
+        input_spec=my_input_spec,
+        inpA="inpA",
+        inpB=1,
+    )
+
+    assert shelly.cmdline == f"executable inpA -o {str(shelly.output_dir / 'inpA_odd')}"
+
+
 def test_shell_cmd_inputs_template_1_st():
     """additional inputs, one uses output_file_template (and argstr)
     testing cmdline when splitter defined