added shell_task decorator

tclose · tclose · commit 83c06e3d0f3c · 2023-05-15T10:27:00.000+10:00
diff --git a/pydra/mark/__init__.py b/pydra/mark/__init__.py
@@ -1,4 +1,4 @@
 from .functions import annotate, task
-from .shell_commands import cmd_arg, cmd_out
+from .shell_commands import shell_task, shell_arg, shell_out
 
-__all__ = ("annotate", "task", "cmd_arg", "cmd_out")
+__all__ = ("annotate", "task", "shell_task", "shell_arg", "shell_out")
diff --git a/pydra/mark/shell_commands.py b/pydra/mark/shell_commands.py
@@ -2,9 +2,136 @@
 from __future__ import annotations
 import typing as ty
 import attrs
+import pydra.engine.specs
 
 
-def cmd_arg(
+def shell_task(
+    klass_or_name: ty.Union[type, str],
+    executable: ty.Optional[str] = None,
+    input_fields: ty.Optional[dict[str, dict]] = None,
+    output_fields: ty.Optional[dict[str, dict]] = None,
+    bases: ty.Optional[list[type]] = None,
+    input_bases: ty.Optional[list[type]] = None,
+    output_bases: ty.Optional[list[type]] = None,
+) -> type:
+    """
+    Construct an analysis class and validate all the components fit together
+
+    Parameters
+    ----------
+    klass_or_name : type or str
+        Either the class decorated by the @shell_task decorator or the name for a
+        dynamically generated class
+    executable : str, optional
+        If dynamically constructing a class (instead of decorating an existing one) the
+        name of the executable to run is provided
+    input_fields : dict[str, dict], optional
+        If dynamically constructing a class (instead of decorating an existing one) the
+        input fields can be provided as a dictionary of dictionaries, where the keys
+        are the name of the fields and the dictionary contents are passed as keyword
+        args to cmd_arg, with the exception of "type", which is used as the type annotation
+        of the field.
+    output_fields : dict[str, dict], optional
+        If dynamically constructing a class (instead of decorating an existing one) the
+        output fields can be provided as a dictionary of dictionaries, where the keys
+        are the name of the fields and the dictionary contents are passed as keyword
+        args to cmd_out, with the exception of "type", which is used as the type annotation
+        of the field.
+    bases : list[type]
+        Base classes for dynamically constructed shell command classes
+    input_bases : list[type]
+        Base classes for the input spec of dynamically constructed shell command classes
+    output_bases : list[type]
+        Base classes for the input spec of dynamically constructed shell command classes
+
+    Returns
+    -------
+    type
+        the shell command task class
+    """
+
+    if isinstance(klass_or_name, str):
+        if None in (executable, input_fields):
+            raise RuntimeError(
+                "Dynamically constructed shell tasks require an executable and "
+                "input_field arguments"
+            )
+        name = klass_or_name
+        if output_fields is None:
+            output_fields = {}
+        if bases is None:
+            bases = [pydra.engine.task.ShellCommandTask]
+        if input_bases is None:
+            input_bases = [pydra.engine.specs.ShellSpec]
+        if output_bases is None:
+            output_bases = [pydra.engine.specs.ShellOutSpec]
+        Inputs = type("Inputs", tuple(input_bases), input_fields)
+        Outputs = type("Outputs", tuple(output_bases), output_fields)
+    else:
+        if (
+            executable,
+            input_fields,
+            output_fields,
+            bases,
+            input_bases,
+            output_bases,
+        ) != (None, None, None, None, None, None):
+            raise RuntimeError(
+                "When used as a decorator on a class `shell_task` should not be provided "
+                "executable, input_field or output_field arguments"
+            )
+        klass = klass_or_name
+        name = klass.__name__
+        try:
+            executable = klass.executable
+        except KeyError:
+            raise RuntimeError(
+                "Classes decorated by `shell_task` should contain an `executable` attribute "
+                "specifying the shell tool to run"
+            )
+        try:
+            Inputs = klass.Inputs
+        except KeyError:
+            raise RuntimeError(
+                "Classes decorated by `shell_task` should contain an `Inputs` class attribute "
+                "specifying the inputs to the shell tool"
+            )
+        if not issubclass(Inputs, pydra.engine.specs.ShellSpec):
+            Inputs = type("Inputs", (Inputs, pydra.engine.specs.ShellSpec), {})
+        try:
+            Outputs = klass.Outputs
+        except KeyError:
+            Outputs = type("Outputs", (pydra.engine.specs.ShellOutSpec,))
+        bases = [klass]
+        if not issubclass(klass, pydra.engine.task.ShellCommandTask):
+            bases.append(pydra.engine.task.ShellCommandTask)
+
+    Inputs = attrs.define(kw_only=True, slots=False)(Inputs)
+    Outputs = attrs.define(kw_only=True, slots=False)(Outputs)
+
+    dct = {
+        "executable": executable,
+        "Inputs": Outputs,
+        "Outputs": Inputs,
+        "inputs": attrs.field(factory=Inputs),
+        "outputs": attrs.field(factory=Outputs),
+        "__annotations__": {
+            "executable": str,
+            "inputs": Inputs,
+            "outputs": Outputs,
+        },
+    }
+
+    return attrs.define(kw_only=True, slots=False)(
+        type(
+            name,
+            tuple(bases),
+            dct,
+        )
+    )
+
+
+def shell_arg(
     help_string: str,
     default: ty.Any = attrs.NOTHING,
     argstr: str = None,
@@ -103,7 +230,7 @@ def cmd_arg(
     )
 
 
-def cmd_out(
+def shell_out(
     help_string: str,
     mandatory: bool = False,
     output_file_template: str = None,
diff --git a/pydra/mark/tests/test_shell_commands.py b/pydra/mark/tests/test_shell_commands.py
@@ -3,36 +3,36 @@
 from pathlib import Path
 import attrs
 import pydra.engine
-from pydra.mark import cmd_arg, cmd_out
+from pydra.mark import shell_task, shell_arg, shell_out
 
 
-def test_shell_cmd():
+def test_shell_task_full():
     @attrs.define(kw_only=True, slots=False)
     class LsInputSpec(pydra.specs.ShellSpec):
-        directory: os.PathLike = cmd_arg(
+        directory: os.PathLike = shell_arg(
             help_string="the directory to list the contents of",
             argstr="",
             mandatory=True,
         )
-        hidden: bool = cmd_arg(help_string=("display hidden FS objects"), argstr="-a")
-        long_format: bool = cmd_arg(
+        hidden: bool = shell_arg(help_string=("display hidden FS objects"), argstr="-a")
+        long_format: bool = shell_arg(
             help_string=(
                 "display properties of FS object, such as permissions, size and timestamps "
             ),
             argstr="-l",
         )
-        human_readable: bool = cmd_arg(
+        human_readable: bool = shell_arg(
             help_string="display file sizes in human readable form",
             argstr="-h",
             requires=["long_format"],
         )
-        complete_date: bool = cmd_arg(
+        complete_date: bool = shell_arg(
             help_string="Show complete date in long format",
             argstr="-T",
             requires=["long_format"],
             xor=["date_format_str"],
         )
-        date_format_str: str = cmd_arg(
+        date_format_str: str = shell_arg(
             help_string="format string for ",
             argstr="-D",
             requires=["long_format"],
@@ -44,7 +44,7 @@ def list_outputs(stdout):
 
     @attrs.define(kw_only=True, slots=False)
     class LsOutputSpec(pydra.specs.ShellOutSpec):
-        entries: list = cmd_out(
+        entries: list = shell_out(
             help_string="list of entries returned by ls command", callable=list_outputs
         )
 
