Get timestamps with context manager

Author: ekouts

docs/tutorial_basic.rst

@@ -178,14 +178,14 @@ If everything is configured correctly for your system, you should get an output
 .. code-block:: none
 
   [ReFrame Setup]
-   version:           3.0-dev7 (rev: 85ca676f)
-   command:           './bin/reframe -C tutorial/config/settings.py -c tutorial/example1.py -r'
-   launched by:       user@daint104
-   working directory: '/path/to/reframe'
-   settings file:     'tutorial/config/settings.py'
-   check search path: (R) '/path/to/reframe/tutorial/example1.py'
-   stage directory:   '/path/to/reframe/stage'
-   output directory:  '/path/to/reframe/output'
+    version:           3.0-dev7 (rev: 85ca676f)
+    command:           './bin/reframe -C tutorial/config/settings.py -c tutorial/example1.py -r'
+    launched by:       user@daint104
+    working directory: '/path/to/reframe'
+    settings file:     'tutorial/config/settings.py'
+    check search path: (R) '/path/to/reframe/tutorial/example1.py'
+    stage directory:   '/path/to/reframe/stage'
+    output directory:  '/path/to/reframe/output'
 
   [==========] Running 1 check(s)
   [==========] Started on Wed Jun  3 08:50:46 2020

reframe/frontend/executors/__init__.py

@@ -207,25 +207,34 @@ def _notify_listeners(self, callback_name):
             callback(self)
 
     def _safe_call(self, fn, *args, **kwargs):
+        class update_timestamps:
+            '''Context manager to set the start and finish timestamps.'''
+
+            def __init__(self, obj):
+                self.obj = obj
+
+            def __enter__(self):
+                if fn.__name__ != 'poll':
+                    cs = self.obj._current_stage
+                    self.obj._timestamps[f'{cs}_start'] = time.time()
+
+            def __exit__(self, exc_type, exc_value, traceback):
+                cs = self.obj._current_stage
+                self.obj._timestamps[f'{cs}_finish'] = time.time()
+                self.obj._timestamps['pipeline_end'] = time.time()
+
         if fn.__name__ != 'poll':
             self._current_stage = fn.__name__
-            self._timestamps[f'{self._current_stage}_start'] = time.time()
 
         try:
             with logging.logging_context(self.check) as logger:
-                logger.debug('entering stage: %s' % self._current_stage)
-                ret = fn(*args, **kwargs)
-                self._timestamps[f'{self._current_stage}_finish'] = time.time()
-                self._timestamps['pipeline_end'] = time.time()
-                return ret
+                logger.debug(f'entering stage: {self._current_stage}')
+                with update_timestamps(self):
+                    return fn(*args, **kwargs)
         except ABORT_REASONS:
-            self._timestamps[f'{self._current_stage}_finish'] = time.time()
-            self._timestamps['pipeline_end'] = time.time()
             self.fail()
             raise
         except BaseException as e:
-            self._timestamps[f'{self._current_stage}_finish'] = time.time()
-            self._timestamps['pipeline_end'] = time.time()
             self.fail()
             raise TaskExit from e