moved hooks arguments into single TaskHooks object in run methods

Author: tclose

pydra/engine/core.py

@@ -29,7 +29,7 @@
 from .specs import (
     RuntimeSpec,
     Result,
-    TaskHook,
+    TaskHooks,
 )
 from .helpers import (
     attrs_fields,
@@ -100,10 +100,7 @@ def __init__(
         name: str,
         environment: "Environment | None" = None,
         state_index: "state.StateIndex | None" = None,
-        pre_run: ty.Callable["Task", None] | None = None,
-        pre_run_task: ty.Callable["Task", None] | None = None,
-        post_run_task: ty.Callable["Task", None] | None = None,
-        post_run: ty.Callable["Task", None] | None = None,
+        hooks: TaskHooks | None = None,
     ):
         """
         Initialize a task.
@@ -146,12 +143,7 @@ def __init__(
         self.allow_cache_override = True
         self._checksum = None
         self._uid = uuid4().hex
-        self.hooks = TaskHook(
-            pre_run=pre_run,
-            post_run=post_run,
-            pre_run_task=pre_run_task,
-            post_run_task=post_run_task,
-        )
+        self.hooks = hooks if hooks is not None else TaskHooks()
         self._errored = False
         self._lzout = None
 

pydra/engine/specs.py

@@ -138,6 +138,32 @@ def __eq__(self, other: ty.Any) -> bool:
 OutputsType = ty.TypeVar("OutputType", bound=TaskOutputs)
 
 
+def donothing(*args: ty.Any, **kwargs: ty.Any) -> None:
+    return None
+
+
+@attrs.define(kw_only=True)
+class TaskHooks:
+    """Callable task hooks."""
+
+    pre_run_task: ty.Callable = attrs.field(
+        default=donothing, converter=default_if_none(donothing)
+    )
+    post_run_task: ty.Callable = attrs.field(
+        default=donothing, converter=default_if_none(donothing)
+    )
+    pre_run: ty.Callable = attrs.field(
+        default=donothing, converter=default_if_none(donothing)
+    )
+    post_run: ty.Callable = attrs.field(
+        default=donothing, converter=default_if_none(donothing)
+    )
+
+    def reset(self):
+        for val in ["pre_run_task", "post_run_task", "pre_run", "post_run"]:
+            setattr(self, val, donothing)
+
+
 @attrs.define(kw_only=True, auto_attribs=False, eq=False)
 class TaskDef(ty.Generic[OutputsType]):
     """Base class for all task definitions"""
@@ -161,10 +187,7 @@ def __call__(
         messengers: ty.Iterable[Messenger] | None = None,
         messenger_args: dict[str, ty.Any] | None = None,
         name: str | None = None,
-        pre_run: ty.Callable["Task", None] | None = None,
-        post_run: ty.Callable["Task", None] | None = None,
-        pre_run_task: ty.Callable["Task", None] | None = None,
-        post_run_task: ty.Callable["Task", None] | None = None,
+        hooks: TaskHooks | None = None,
         **kwargs: ty.Any,
     ) -> OutputsType:
         """Create a task from this definition and execute it to produce a result.
@@ -220,10 +243,7 @@ def __call__(
                 result = sub(
                     self,
                     name=name,
-                    pre_run=pre_run,
-                    post_run=post_run,
-                    pre_run_task=pre_run_task,
-                    post_run_task=post_run_task,
+                    hooks=hooks,
                 )
         except TypeError as e:
             # Catch any inadvertent passing of task definition parameters to the
@@ -1254,32 +1274,6 @@ def _generated_output_names(self, stdout: str, stderr: str):
     DEFAULT_COPY_COLLATION = FileSet.CopyCollation.adjacent
 
 
-def donothing(*args: ty.Any, **kwargs: ty.Any) -> None:
-    return None
-
-
-@attrs.define(kw_only=True)
-class TaskHook:
-    """Callable task hooks."""
-
-    pre_run_task: ty.Callable = attrs.field(
-        default=donothing, converter=default_if_none(donothing)
-    )
-    post_run_task: ty.Callable = attrs.field(
-        default=donothing, converter=default_if_none(donothing)
-    )
-    pre_run: ty.Callable = attrs.field(
-        default=donothing, converter=default_if_none(donothing)
-    )
-    post_run: ty.Callable = attrs.field(
-        default=donothing, converter=default_if_none(donothing)
-    )
-
-    def reset(self):
-        for val in ["pre_run_task", "post_run_task", "pre_run", "post_run"]:
-            setattr(self, val, donothing)
-
-
 def split_cmd(cmd: str | None):
     """Splits a shell command line into separate arguments respecting quotes
 

pydra/engine/submitter.py

@@ -29,7 +29,7 @@
 
 if ty.TYPE_CHECKING:
     from .node import Node
-    from .specs import TaskDef, WorkflowDef
+    from .specs import TaskDef, WorkflowDef, TaskHooks
     from .environments import Environment
     from .state import State
 
@@ -169,10 +169,7 @@ def __call__(
         self,
         task_def: "TaskDef",
         name: str | None = "task",
-        pre_run: ty.Callable["Task", None] | None = None,
-        post_run: ty.Callable["Task", None] | None = None,
-        pre_run_task: ty.Callable["Task", None] | None = None,
-        post_run_task: ty.Callable["Task", None] | None = None,
+        hooks: "TaskHooks | None" = None,
     ):
         """Submitter run function."""
 
@@ -203,10 +200,7 @@ def Split(defn: TaskDef, output_types: dict):
             submitter=self,
             name=name,
             environment=self.environment,
-            pre_run=pre_run,
-            post_run=post_run,
-            pre_run_task=pre_run_task,
-            post_run_task=post_run_task,
+            hooks=hooks,
         )
         try:
             self.run_start_time = datetime.now()

pydra/engine/tests/test_task.py

@@ -11,11 +11,10 @@
 from pydra.design import python, shell, workflow
 from pydra.utils.messenger import FileMessenger, PrintMessenger, collect_messages
 from ..task import AuditFlag
-from pydra.engine.specs import argstr_formatting, ShellDef, ShellOutputs
+from pydra.engine.specs import argstr_formatting, ShellDef, ShellOutputs, TaskHooks
 from pydra.engine.helpers import list_fields, print_help
 from pydra.engine.submitter import Submitter
 from pydra.engine.core import Task
-from .utils import BasicWorkflow
 from pydra.utils import default_run_cache_dir
 from pydra.utils.typing import (
     MultiInputObj,
@@ -1209,7 +1208,7 @@ def test_taskhooks_1(tmpdir: Path, capsys):
     def myhook(task, *args):
         print("I was called")
 
-    FunAddTwo(a=1)(cache_dir=cache_dir, pre_run=myhook)
+    FunAddTwo(a=1)(cache_dir=cache_dir, hooks=TaskHooks(pre_run=myhook))
     captured = capsys.readouterr()
     assert "I was called\n" in captured.out
     del captured
@@ -1231,10 +1230,12 @@ def myhook(task, *args):
     # set all hooks
     FunAddTwo(a=1)(
         cache_dir=cache_dir,
-        pre_run=myhook,
-        post_run=myhook,
-        pre_run_task=myhook,
-        post_run_task=myhook,
+        hooks=TaskHooks(
+            pre_run=myhook,
+            post_run=myhook,
+            pre_run_task=myhook,
+            post_run_task=myhook,
+        ),
     )
     captured = capsys.readouterr()
     assert captured.out.count("I was called\n") == 4
@@ -1258,10 +1259,12 @@ def myhook_postrun(task, *args):
 
     FunAddTwo(a=1)(
         cache_dir=tmpdir,
-        pre_run=myhook_prerun,
-        post_run=myhook_postrun,
-        pre_run_task=myhook_prerun_task,
-        post_run_task=myhook_postrun_task,
+        hooks=TaskHooks(
+            pre_run=myhook_prerun,
+            post_run=myhook_postrun,
+            pre_run_task=myhook_prerun_task,
+            post_run_task=myhook_postrun_task,
+        ),
     )
 
     captured = capsys.readouterr()
@@ -1307,7 +1310,8 @@ def myhook_postrun(task, result, *args):
 
     with pytest.raises(Exception):
         FunAddTwo(a="one")(
-            cache_dir=tmpdir, post_run=myhook_postrun, post_run_task=myhook_postrun_task
+            cache_dir=tmpdir,
+            hooks=TaskHooks(post_run=myhook_postrun, post_run_task=myhook_postrun_task),
         )
 
     captured = capsys.readouterr()