Merge pull request #389 from chasejohnson3/master

Added allowed types for an output_spec field

Author: satra

pydra/engine/core.py

@@ -486,7 +486,9 @@ def _collect_outputs(self, output_dir):
         run_output = self.output_
         output_klass = make_klass(self.output_spec)
         output = output_klass(**{f.name: None for f in attr.fields(output_klass)})
-        other_output = output.collect_additional_outputs(self.inputs, output_dir)
+        other_output = output.collect_additional_outputs(
+            self.inputs, output_dir, run_output
+        )
         return attr.evolve(output, **run_output, **other_output)
 
     def split(self, splitter, overwrite=False, **kwargs):

pydra/engine/specs.py

@@ -95,7 +95,7 @@ def __setattr__(self, name, value):
             # validate all fields that have set a validator
             attr.validate(self)
 
-    def collect_additional_outputs(self, inputs, output_dir):
+    def collect_additional_outputs(self, inputs, output_dir, outputs):
         """Get additional outputs."""
         return {}
 
@@ -439,12 +439,21 @@ class ShellOutSpec:
     stderr: ty.Union[File, str]
     """The process' standard input."""
 
-    def collect_additional_outputs(self, inputs, output_dir):
+    def collect_additional_outputs(self, inputs, output_dir, outputs):
         """Collect additional outputs from shelltask output_spec."""
         additional_out = {}
         for fld in attr_fields(self):
             if fld.name not in ["return_code", "stdout", "stderr"]:
-                if fld.type in [File, MultiOutputFile]:
+                if fld.type in [
+                    File,
+                    MultiOutputFile,
+                    Directory,
+                    int,
+                    float,
+                    bool,
+                    str,
+                    list,
+                ]:
                     # assuming that field should have either default or metadata, but not both
                     if (
                         fld.default is None or fld.default == attr.NOTHING
@@ -457,9 +466,21 @@ def collect_additional_outputs(self, inputs, output_dir):
                             fld, output_dir
                         )
                     elif fld.metadata:
-                        additional_out[fld.name] = self._field_metadata(
-                            fld, inputs, output_dir
-                        )
+                        if (
+                            fld.type in [int, float, bool, str, list]
+                            and "callable" not in fld.metadata
+                        ):
+                            raise AttributeError(
+                                f"{fld.type} has to have a callable in metadata"
+                            )
+                        else:
+                            additional_out[fld.name] = self._field_metadata(
+                                fld, inputs, output_dir, outputs
+                            )
+                #                        else:
+                #                            additional_out[fld.name] = self._field_metadata(
+                #                                fld, inputs, output_dir, outputs
+                #                            )
                 else:
                     raise Exception("not implemented (collect_additional_output)")
         return additional_out
@@ -486,7 +507,7 @@ def generated_output_names(self, inputs, output_dir):
                         output_names.append(fld.name)
                     elif (
                         fld.metadata
-                        and self._field_metadata(fld, inputs, output_dir)
+                        and self._field_metadata(fld, inputs, output_dir, outputs=None)
                         != attr.NOTHING
                     ):
                         output_names.append(fld.name)
@@ -522,7 +543,7 @@ def _field_defaultvalue(self, fld, output_dir):
             else:
                 raise AttributeError(f"no file matches {default.name}")
 
-    def _field_metadata(self, fld, inputs, output_dir):
+    def _field_metadata(self, fld, inputs, output_dir, outputs=None):
         """Collect output file if metadata specified."""
         if self._check_requires(fld, inputs) is False:
             return attr.NOTHING
@@ -551,6 +572,10 @@ def _field_metadata(self, fld, inputs, output_dir):
                     call_args_val[argnm] = output_dir
                 elif argnm == "inputs":
                     call_args_val[argnm] = inputs
+                elif argnm == "stdout":
+                    call_args_val[argnm] = outputs["stdout"]
+                elif argnm == "stderr":
+                    call_args_val[argnm] = outputs["stderr"]
                 else:
                     try:
                         call_args_val[argnm] = getattr(inputs, argnm)

pydra/engine/tests/test_shelltask.py

@@ -3,7 +3,7 @@
 import os, sys
 import pytest
 from pathlib import Path
-
+import re
 
 from ..task import ShellCommandTask
 from ..submitter import Submitter
@@ -2837,6 +2837,96 @@ def test_shell_cmd_outputspec_6a(tmpdir, plugin, results_function):
     assert res.output.new_files.exists()
 
 
+@pytest.mark.parametrize("results_function", [result_no_submitter, result_submitter])
+def test_shell_cmd_outputspec_7a(tmpdir, plugin, results_function):
+    """
+        customised output_spec, adding int and str to the output,
+        requiring two callables with parameters stdout and stderr
+    """
+    cmd = "echo"
+    args = ["newfile_1.txt", "newfile_2.txt"]
+
+    def get_file_index(stdout):
+        stdout = re.sub(r".*_", "", stdout)
+        stdout = re.sub(r".txt", "", stdout)
+        print(stdout)
+        return int(stdout)
+
+    def get_stderr(stderr):
+        return f"stderr: {stderr}"
+
+    my_output_spec = SpecInfo(
+        name="Output",
+        fields=[
+            (
+                "out1",
+                attr.ib(
+                    type=File,
+                    metadata={
+                        "output_file_template": "{args}",
+                        "help_string": "output file",
+                    },
+                ),
+            ),
+            (
+                "out_file_index",
+                attr.ib(
+                    type=int,
+                    metadata={"help_string": "output file", "callable": get_file_index},
+                ),
+            ),
+            (
+                "stderr_field",
+                attr.ib(
+                    type=str,
+                    metadata={
+                        "help_string": "The standard error output",
+                        "callable": get_stderr,
+                    },
+                ),
+            ),
+        ],
+        bases=(ShellOutSpec,),
+    )
+
+    shelly = ShellCommandTask(
+        name="shelly", executable=cmd, args=args, output_spec=my_output_spec
+    ).split("args")
+
+    results = results_function(shelly, plugin)
+    for index, res in enumerate(results):
+        assert res.output.out_file_index == index + 1
+        assert res.output.stderr_field == f"stderr: {res.output.stderr}"
+
+
+def test_shell_cmd_outputspec_7b_error():
+    """
+        customised output_spec, adding Int to the output,
+        requiring a function to collect output
+    """
+    cmd = "echo"
+    args = ["newfile_1.txt", "newfile_2.txt"]
+
+    my_output_spec = SpecInfo(
+        name="Output",
+        fields=[
+            (
+                "out",
+                attr.ib(
+                    type=int, metadata={"help_string": "output file", "value": "val"}
+                ),
+            )
+        ],
+        bases=(ShellOutSpec,),
+    )
+    shelly = ShellCommandTask(
+        name="shelly", executable=cmd, args=args, output_spec=my_output_spec
+    ).split("args")
+    with pytest.raises(Exception) as e:
+        shelly()
+    assert "has to have a callable" in str(e.value)
+
+
 @pytest.mark.parametrize("results_function", [result_no_submitter, result_submitter])
 def test_shell_cmd_state_outputspec_1(plugin, results_function, tmpdir):
     """