[feat] Re-implement pipeline hook mechanism and add a post-init hook

docs/regression_test_api.rst

@@ -29,6 +29,23 @@ Regression Test Class Decorators
 Pipeline Hooks
 --------------
 
+.. versionadded:: 2.20
+
+
+Pipeline hooks is an easy way to perform operations while the test traverses the execution pipeline.
+You can attach arbitrary functions to run before or after any pipeline stage, which are called *pipeline hooks*.
+Multiple hooks can be attached before or after the same pipeline stage, in which case the order of execution will match the order in which the functions are defined in the class body of the test.
+A single hook can also be applied to multiple stages and it will be executed multiple times.
+All pipeline hooks of a test class are inherited by its subclasses.
+Subclasses may override a pipeline hook of their parents by redefining the hook function and re-attaching it at the same pipeline stage.
+There are seven pipeline stages where you can attach test methods: ``init``, ``setup``, ``compile``, ``run``, ``sanity``, ``performance`` and ``cleanup``.
+The ``init`` stage is not a real pipeline stage, but it refers to the test initialization.
+
+Hooks attached to any stage will run exactly before or after this stage executes.
+So although a "post-init" and a "pre-setup" hook will both run *after* a test has been initialized and *before* the test goes through the first pipeline stage, they will execute in different times:
+the post-init hook will execute *right after* the test is initialized.
+The framework will then continue with other activities and it will execute the pre-setup hook *just before* it schedules the test for executing its setup stage.
+
 .. autodecorator:: reframe.core.decorators.run_after(stage)
 
 .. autodecorator:: reframe.core.decorators.run_before(stage)

reframe/core/decorators.py

@@ -214,6 +214,13 @@ def _fn(*args, **kwargs):
     return deco
 
 
+# Valid pipeline stages that users can specify in the `run_before()` and
+# `run_after()` decorators
+_USER_PIPELINE_STAGES = (
+    'init', 'setup', 'compile', 'run', 'sanity', 'performance', 'cleanup'
+)
+
+
 def run_before(stage):
     '''Decorator for attaching a test method to a pipeline stage.
 
@@ -226,19 +233,57 @@ def run_before(stage):
     The ``stage`` argument can be any of ``'setup'``, ``'compile'``,
     ``'run'``, ``'sanity'``, ``'performance'`` or ``'cleanup'``.
 
-    .. versionadded:: 2.20
     '''
+    if stage not in _USER_PIPELINE_STAGES:
+        raise ValueError(f'invalid pipeline stage specified: {stage!r}')
+
+    if stage == 'init':
+        raise ValueError('pre-init hooks are not allowed')
+
     return _runx('pre_' + stage)
 
 
 def run_after(stage):
     '''Decorator for attaching a test method to a pipeline stage.
 
-    This is completely analogous to the
-    :py:attr:`reframe.core.decorators.run_before`.
+    This is analogous to the :py:attr:`~reframe.core.decorators.run_before`,
+    except that ``'init'`` can also be used as the ``stage`` argument. In this
+    case, the hook will execute right after the test is initialized (i.e.
+    after the :func:`__init__` method is called), before entering the test's
+    pipeline. In essence, a post-init hook is equivalent to defining
+    additional :func:`__init__` functions in the test. All the other
+    properties of pipeline hooks apply equally here. The following code
+
+    .. code-block:: python
+
+       @rfm.run_after('init')
+       def foo(self):
+           self.x = 1
+
+
+    is equivalent to
+
+    .. code-block:: python
+
+       def __init__(self):
+           self.x = 1
+
+    .. versionchanged:: 3.5.2
+       Add the ability to define post-init hooks in tests.
 
-    .. versionadded:: 2.20
     '''
+
+    if stage not in _USER_PIPELINE_STAGES:
+        raise ValueError(f'invalid pipeline stage specified: {stage!r}')
+
+    # Map user stage names to the actual pipeline functions if needed
+    if stage == 'init':
+        stage = '__init__'
+    elif stage == 'compile':
+        stage = 'compile_wait'
+    elif stage == 'run':
+        stage = 'run_wait'
+
     return _runx('post_' + stage)
 
 

reframe/core/hooks.py

@@ -0,0 +1,135 @@
+# Copyright 2016-2021 Swiss National Supercomputing Centre (CSCS/ETH Zurich)
+# ReFrame Project Developers. See the top-level LICENSE file for details.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+import contextlib
+import functools
+
+import reframe.utility as util
+
+
+def attach_hooks(hooks):
+    '''Attach pipeline hooks to phase ``name''.
+
+    This function returns a decorator for pipeline functions that will run the
+    registered hooks before and after the function.
+
+    If ``name'' is :class:`None`, both pre- and post-hooks will run, otherwise
+    only the hooks of the phase ``name'' will be executed.
+    '''
+
+    def _deco(func):
+        def select_hooks(obj, kind):
+            phase = kind + func.__name__
+            if phase not in hooks:
+                return []
+
+            return [h for h in hooks[phase]
+                    if h.__name__ not in obj._disabled_hooks]
+
+        @functools.wraps(func)
+        def _fn(obj, *args, **kwargs):
+            for h in select_hooks(obj, 'pre_'):
+                h(obj)
+
+            func(obj, *args, **kwargs)
+            for h in select_hooks(obj, 'post_'):
+                h(obj)
+
+        return _fn
+
+    return _deco
+
+
+class Hook:
+    '''A pipeline hook.
+
+    This is essentially a function wrapper that hashes the functions by name,
+    since we want hooks to be overriden by name in subclasses.
+    '''
+
+    def __init__(self, fn):
+        self.__fn = fn
+
+    def __getattr__(self, attr):
+        return getattr(self.__fn, attr)
+
+    @property
+    def fn(self):
+        return self.__fn
+
+    def __hash__(self):
+        return hash(self.__name__)
+
+    def __eq__(self, other):
+        if not isinstance(other, type(self)):
+            return NotImplemented
+
+        return self.__name__ == other.__name__
+
+    def __call__(self, *args, **kwargs):
+        return self.__fn(*args, **kwargs)
+
+    def __repr__(self):
+        return repr(self.__fn)
+
+
+class HookRegistry:
+    '''Global hook registry.'''
+
+    @classmethod
+    def create(cls, namespace):
+        '''Create a hook registry from a class namespace.
+
+        Hook functions have an `_rfm_attach` attribute that specify the stages
+        of the pipeline where they must be attached. Dependencies will be
+        resolved first in the post-setup phase if not assigned elsewhere.
+        '''
+
+        local_hooks = {}
+        fn_with_deps = []
+        for v in namespace.values():
+            if hasattr(v, '_rfm_attach'):
+                for phase in v._rfm_attach:
+                    try:
+                        local_hooks[phase].append(Hook(v))
+                    except KeyError:
+                        local_hooks[phase] = [Hook(v)]
+
+            with contextlib.suppress(AttributeError):
+                if v._rfm_resolve_deps:
+                    fn_with_deps.append(Hook(v))
+
+        if fn_with_deps:
+            local_hooks['post_setup'] = (
+                fn_with_deps + local_hooks.get('post_setup', [])
+            )
+
+        return cls(local_hooks)
+
+    def __init__(self, hooks=None):
+        self.__hooks = {}
+        if hooks is not None:
+            self.update(hooks)
+
+    def __getitem__(self, key):
+        return self.__hooks[key]
+
+    def __setitem__(self, key, name):
+        self.__hooks[key] = name
+
+    def __contains__(self, key):
+        return key in self.__hooks
+
+    def __getattr__(self, name):
+        return getattr(self.__hooks, name)
+
+    def update(self, hooks):
+        for phase, hks in hooks.items():
+            self.__hooks.setdefault(phase, util.OrderedSet())
+            for h in hks:
+                self.__hooks[phase].add(h)
+
+    def __repr__(self):
+        return repr(self.__hooks)

reframe/core/meta.py

@@ -8,11 +8,13 @@
 #
 
 
-from reframe.core.exceptions import ReframeSyntaxError
 import reframe.core.namespaces as namespaces
 import reframe.core.parameters as parameters
 import reframe.core.variables as variables
 
+from reframe.core.exceptions import ReframeSyntaxError
+from reframe.core.hooks import HookRegistry
+
 
 class RegressionTestMeta(type):
 
@@ -143,27 +145,12 @@ def __init__(cls, name, bases, namespace, **kwargs):
         # Set up the hooks for the pipeline stages based on the _rfm_attach
         # attribute; all dependencies will be resolved first in the post-setup
         # phase if not assigned elsewhere
-        hooks = {}
-        fn_with_deps = []
-        for v in namespace.values():
-            if hasattr(v, '_rfm_attach'):
-                for phase in v._rfm_attach:
-                    try:
-                        hooks[phase].append(v)
-                    except KeyError:
-                        hooks[phase] = [v]
-
-            try:
-                if v._rfm_resolve_deps:
-                    fn_with_deps.append(v)
-            except AttributeError:
-                pass
-
-        if fn_with_deps:
-            hooks['post_setup'] = fn_with_deps + hooks.get('post_setup', [])
+        hooks = HookRegistry.create(namespace)
+        for b in bases:
+            if hasattr(b, '_rfm_pipeline_hooks'):
+                hooks.update(getattr(b, '_rfm_pipeline_hooks'))
 
-        cls._rfm_pipeline_hooks = hooks
-        cls._rfm_disabled_hooks = set()
+        cls._rfm_pipeline_hooks = hooks  # HookRegistry(local_hooks)
         cls._final_methods = {v.__name__ for v in namespace.values()
                               if hasattr(v, '_rfm_final')}