Add log prefix context manager to make prefix cleanup/restoration easier (#635).

borgmatic/commands/borgmatic.py

@@ -41,9 +41,9 @@ from borgmatic.hooks import command, dispatch
 from borgmatic.hooks.monitoring import monitor
 from borgmatic.logger import (
     DISABLED,
+    Log_prefix,
     add_custom_log_levels,
     configure_logging,
-    set_log_prefix,
     should_do_markup,
 )
 from borgmatic.signals import configure_signals
@@ -134,10 +134,10 @@ def run_configuration(config_filename, config, config_paths, arguments):
                 (repo, 0),
             )
 
-        try:
-            while not repo_queue.empty():
-                repository, retry_num = repo_queue.get()
-                set_log_prefix(repository.get('label', repository['path']))
+        while not repo_queue.empty():
+            repository, retry_num = repo_queue.get()
+
+            with Log_prefix(repository.get('label', repository['path'])):
                 logger.debug('Running actions for repository')
                 timeout = retry_num * retry_wait
                 if timeout:
@@ -179,8 +179,6 @@ def run_configuration(config_filename, config, config_paths, arguments):
                     )
                     encountered_error = error
                     error_repository = repository['path']
-        finally:
-            set_log_prefix(config_filename)
 
     try:
         if monitoring_hooks_are_activated:
@@ -817,9 +815,8 @@ def collect_configuration_run_summary_logs(configs, config_paths, arguments):
     # Execute the actions corresponding to each configuration file.
     json_results = []
 
-    try:
-        for config_filename, config in configs.items():
-            set_log_prefix(config_filename)
+    for config_filename, config in configs.items():
+        with Log_prefix(config_filename):
             results = list(run_configuration(config_filename, config, config_paths, arguments))
             error_logs = tuple(
                 result for result in results if isinstance(result, logging.LogRecord)
@@ -838,8 +835,6 @@ def collect_configuration_run_summary_logs(configs, config_paths, arguments):
                 )
                 if results:
                     json_results.extend(results)
-    finally:
-        set_log_prefix(None)
 
     if 'umount' in arguments:
         logger.info(f"Unmounting mount point {arguments['umount'].mount_point}")

borgmatic/execute.py

@@ -311,9 +311,7 @@ def execute_command(
     if not run_to_completion:
         return process
 
-    try:
-        original_log_prefix = borgmatic.logger.get_log_prefix()
-        borgmatic.logger.set_log_prefix(None)  # Log command output without any prefix.
+    with borgmatic.logger.Log_prefix(None):  # Log command output without any prefix.
         log_outputs(
             (process,),
             (input_file, output_file),
@@ -321,8 +319,6 @@ def execute_command(
             borg_local_path,
             borg_exit_codes,
         )
-    finally:
-        borgmatic.logger.set_log_prefix(original_log_prefix)
 
 
 def execute_command_and_capture_output(

borgmatic/logger.py

@@ -186,28 +186,52 @@ def add_custom_log_levels():  # pragma: no cover
     add_logging_level('DISABLED', DISABLED)
 
 
-def get_log_prefix():
+class Log_prefix:
     '''
-    Return the log prefix (as set with set_log_prefix()) for the first log handler. If there is no
-    such prefix, return None.
-    '''
-    for handler in logging.getLogger().handlers:
-        defaults = handler.formatter._style._defaults
-
-        if not defaults:
-            return None
+    A Python context manager for setting a log prefix so that it shows up in every subsequent
+    logging message for the duration of the context manager. For this to work, it relies on each
+    logging formatter to be initialized with "{prefix}" somewhere in its logging format.
 
-        return defaults.get('prefix').rstrip().rstrip(':')
+    Example use as a context manager:
 
+        
+       with borgmatic.logger.Log_prefix('myprefix'):
+            do_something_that_logs()
 
-def set_log_prefix(prefix):
+    For the scope of that "with" statement, any logs created are prefixed with "myprefix: ".
+    Afterwards, the prefix gets restored to whatever it was prior to the context manager.
     '''
-    Given a prefix string, set it onto the formatter defaults for every logging handler so that it
-    shows up in every subsequent logging message. For this to work, this relies on each logging
-    formatter to be initialized with "{prefix}" somewhere in its logging format.
-    '''
-    for handler in logging.getLogger().handlers:
-        handler.formatter._style._defaults = {'prefix': f'{prefix}: ' if prefix else ''}
+    def __init__(self, prefix):
+        '''
+        Given the desired log prefix, save it for use below. Set prefix to None to disable any
+        prefix from getting logged.
+        '''
+        self.prefix = prefix
+        self.original_prefix = None
+
+    def __enter__(self):
+        '''
+        Set the prefix onto the formatter defaults for every logging handler so that the prefix ends
+        up in every log message. But first, save off any original prefix so that it can be restored
+        below.
+        '''
+        try:
+            self.original_prefix = next(
+                handler.formatter._style._defaults.get('prefix').rstrip().rstrip(':')
+                for handler in logging.getLogger().handlers
+            )
+        except (StopIteration, AttributeError):
+            self.original_prefix = None
+
+        for handler in logging.getLogger().handlers:
+            handler.formatter._style._defaults = {'prefix': f'{self.prefix}: ' if self.prefix else ''}
+
+    def __exit__(self, exception, value, traceback):
+        '''
+        Restore any original prefix.
+        '''
+        for handler in logging.getLogger().handlers:
+            handler.formatter._style._defaults = {'prefix': f'{self.original_prefix}: ' if self.original_prefix else ''}
 
 
 def configure_logging(