From e48ff097881cb726181c187d94aad8dee0a3950b Mon Sep 17 00:00:00 2001
From: George Bisbas <g.bisbas18@imperial.ac.uk>
Date: Mon, 15 Apr 2024 17:55:27 +0100
Subject: [PATCH] docs: Simplify implementation

---
 devito/logger.py            | 14 +++++++-------
 devito/operator/operator.py |  9 ++++-----
 devito/parameters.py        | 18 ++++++++++--------
 3 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/devito/logger.py b/devito/logger.py
index fe751321629..fa8b6684076 100644
--- a/devito/logger.py
+++ b/devito/logger.py
@@ -13,13 +13,13 @@
 stream_handler = logging.StreamHandler()
 logger.addHandler(stream_handler)
 
-# Add extra logging levels
-DEBUG = logging.DEBUG  # value=10
+# Add extra logging levels (note: INFO has value=20, WARNING has value=30)
+DEBUG = logging.DEBUG
 PERF = 19
-INFO = logging.INFO  # value=20
-WARNING = logging.WARNING  # value=30
-ERROR = logging.ERROR  # value=40
-CRITICAL = logging.CRITICAL  # value=50
+INFO = logging.INFO
+WARNING = logging.WARNING
+ERROR = logging.ERROR
+CRITICAL = logging.CRITICAL
 
 logging.addLevelName(PERF, "PERF")
 
@@ -77,7 +77,7 @@ def set_log_level(level, comm=None):
     """
     from devito import configuration
 
-    if comm is not None:
+    if comm is not None and configuration['mpi']:
         if comm.rank != 0:
             logger.removeHandler(stream_handler)
             logger.addHandler(logging.NullHandler())
diff --git a/devito/operator/operator.py b/devito/operator/operator.py
index 36d46853443..c4b4626bad9 100644
--- a/devito/operator/operator.py
+++ b/devito/operator/operator.py
@@ -6,7 +6,7 @@
 from cached_property import cached_property
 from sympy import sympify
 
-from devito import mpi_switch_log
+from devito import switch_log_level
 from devito.arch import compiler_registry, platform_registry
 from devito.data import default_allocator
 from devito.exceptions import InvalidOperator, ExecutionError
@@ -873,12 +873,11 @@ def apply(self, **kwargs):
         self._postprocess_arguments(args, **kwargs)
 
         # In case MPI is used restrict result logging to one rank only
-        if configuration['mpi']:
-            with mpi_switch_log(log_level='DEBUG', comm=args.comm):
-                return self._emit_apply_profiling(args)
-        else:
+        with switch_log_level(comm=args.comm):
             return self._emit_apply_profiling(args)
 
+        return self._emit_apply_profiling(args)
+
     # Performance profiling
     def _emit_build_profiling(self):
         if not is_log_enabled_for('PERF'):
diff --git a/devito/parameters.py b/devito/parameters.py
index 3977a835eff..4066f3cc022 100644
--- a/devito/parameters.py
+++ b/devito/parameters.py
@@ -8,7 +8,7 @@
 from devito.tools import Signer, filter_ordered
 
 __all__ = ['configuration', 'init_configuration', 'print_defaults', 'print_state',
-           'switchconfig', 'mpi_switch_log']
+           'switchconfig', 'switch_log_level']
 
 # Be EXTREMELY careful when writing to a Parameters dictionary
 # Read here for reference: http://wiki.c2.com/?GlobalVariablesAreBad
@@ -258,24 +258,26 @@ def wrapper(*args, **kwargs):
         return wrapper
 
 
-class mpi_switch_log(switchconfig):
+class switch_log_level(object):
     """
     A context manager subclassing `switchconfig` to temporarily change
     MPI logging.
     """
 
-    def __init__(self, **params):
-        self.params = {k.replace('_', '-'): v for k, v in params.items()}
-        self.previous = {}
+    def __init__(self, comm):
 
-        comm = self.params.pop('comm')
+        self.level = configuration['log-level']
+        self.comm = comm
 
         # Limit logging to rank 0
-        set_log_level(self.params['log-level'], comm=comm)
+        set_log_level(self.level, comm)
+
+    def __enter__(self):
+        return
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         # Reinstate logging upon exit
-        set_log_level(self.previous['log-level'])
+        set_log_level(self.level)
         logger.addHandler(stream_handler)