log_helper: fallback to root logger

The log_helper was refactored to fallback to the root logger when logging is enabled. Now, instead of each logger having its own level and handler, we reuse the level of the root logger by setting the level of the loggers to logging.NOTSET. We also keep the default propagation behavior so instead of having a custom handler for each logger, they will keep the NullHandler and when logging they will use the handler of the root logger. This also allowed us to set the JSON formatter only for the root logger, since it's the only one actually logging things. The base class for loggers was also removed because the way it was used it would replace the base class for every logger, which would affect modules importing oletools files. We now use an adapter that is only valid for loggers created by the log_helper. The hook for exceptions now wraps any hook previosly set instead of replacing it altogether and this is only done when calling modules as script. Finally, the log_at_current_level() function was replaced by a print_str() function defined in the adapter that will print instead of logging when not using JSON so we can keep the old behavior.

log_helper: fallback to root logger
The log_helper was refactored to fallback to the root logger when logging is enabled. Now, instead of each logger having its own level and handler, we reuse the level of the root logger by setting the level of the loggers to logging.NOTSET. We also keep the default propagation behavior so instead of having a custom handler for each logger, they will keep the NullHandler and when logging they will use the handler of the root logger. This also allowed us to set the JSON formatter only for the root logger, since it's the only one actually logging things. The base class for loggers was also removed because the way it was used it would replace the base class for every logger, which would affect modules importing oletools files. We now use an adapter that is only valid for loggers created by the log_helper. The hook for exceptions now wraps any hook previosly set instead of replacing it altogether and this is only done when calling modules as script. Finally, the log_at_current_level() function was replaced by a print_str() function defined in the adapter that will print instead of logging when not using JSON so we can keep the old behavior.
Samir Aguiar · Christian Herdtweck
1 parent 660ea766
Showing 5 changed files with 103 additions and 85 deletions
oletools/util/log_helper/_json_formatter.py
oletools/util/log_helper/_logger_adapter.py
oletools/util/log_helper/_logger_class.py
oletools/util/log_helper/_root_logger_wrapper.py
oletools/util/log_helper/log_helper.py
@@ -10,18 +10,15 @@ class JsonFormatter(logging.Formatter):
  
     def format(self, record):
         """
-        We accept messages that are either dictionaries or not.
-        When we have dictionaries we can just serialize it as JSON right away.
+        Since we don't buffer messages, we always prepend messages with a comma to make
+        the output JSON-compatible. The only exception is when printing the first line,
+        so we need to keep track of it.
         """
-        trailing_comma = ','
+        json_dict = dict(msg=record.msg, level=record.levelname)
+        formatted_message = '    ' + json.dumps(json_dict)
  
         if self._is_first_line:
-            trailing_comma = ''
             self._is_first_line = False
+            return formatted_message
  
-        json_dict = record.msg \
-            if isinstance(record.msg, dict) \
-            else dict(msg=record.msg)
-        json_dict['level'] = record.levelname
-
-        return trailing_comma + '    ' + json.dumps(json_dict)
+        return ', ' + formatted_message
+import logging
+from . import _root_logger_wrapper
+
+
+class OletoolsLoggerAdapter(logging.LoggerAdapter):
+    """
+    Adapter class for all loggers returned by the logging module.
+    """
+    _json_enabled = None
+
+    def print_str(self, message):
+        """
+        This function replaces normal print() calls so we can format them as JSON
+        when needed or just print them right away otherwise.
+        """
+        if self._json_enabled and self._json_enabled():
+            # Messages from this function should always be printed,
+            # so when using JSON we log using the same level that set
+            self.log(_root_logger_wrapper.level(), message)
+        else:
+            print(message)
+
+    def set_json_enabled_function(self, json_enabled):
+        """
+        Set a function to be called to check whether JSON output is enabled.
+        """
+        self._json_enabled = json_enabled
+
+    def level(self):
+        return self.logger.level
-import logging
-from . import _root_logger_wrapper
-
-
-class OletoolsLogger(logging.Logger):
-    """
-    Default class for all loggers returned by the logging module.
-    """
-    def __init__(self, name, level=logging.NOTSET):
-        super(self.__class__, self).__init__(name, level)
-
-    def log_at_current_level(self, message):
-        """
-        Logs the message using the current level.
-        This is useful for messages that should always appear,
-        such as banners.
-        """
-
-        level = _root_logger_wrapper.get_root_logger_level() \
-            if _root_logger_wrapper.is_logging_initialized() \
-            else self.level
-
-        self.log(level, message)
@@ -3,12 +3,22 @@ import logging
  
 def is_logging_initialized():
     """
-    We use the same strategy as the logging module
-    when checking if the logging was initialized -
-    look for handlers in the root logger
+    We use the same strategy as the logging module when checking if
+    the logging was initialized - look for handlers in the root logger
     """
     return len(logging.root.handlers) > 0
  
  
-def get_root_logger_level():
+def set_formatter(fmt):
+    """
+    Set the formatter to be used by every handler of the root logger.
+    """
+    if not is_logging_initialized():
+        return
+
+    for handler in logging.root.handlers:
+        handler.setFormatter(fmt)
+
+
+def level():
     return logging.root.level
@@ -41,9 +41,8 @@ General logging helpers
 # TODO:
  
  
-from __future__ import print_function
 from ._json_formatter import JsonFormatter
-from ._logger_class import OletoolsLogger
+from ._logger_adapter import OletoolsLoggerAdapter
 from . import _root_logger_wrapper
 import logging
 import sys
@@ -66,8 +65,6 @@ class LogHelper:
         self._all_names = set()  # set so we do not have duplicates
         self._use_json = False
         self._is_enabled = False
-        self._json_formatter = JsonFormatter()
-        logging.setLoggerClass(OletoolsLogger)
  
     def get_or_create_silent_logger(self, name=DEFAULT_LOGGER_NAME, level=logging.CRITICAL + 1):
         """
@@ -75,13 +72,23 @@ class LogHelper:
         as the handler (to avoid printing to the console).
         By default we also use a higher logging level so every message will
         be ignored.
-        This is useful when we don't want to print anything when the logger
-        is not configured by the main application.
+        This will prevent oletools from logging unnecessarily when being imported
+        from external tools.
         """
         return self._get_or_create_logger(name, level, logging.NullHandler())
  
     def enable_logging(self, use_json, level, log_format=DEFAULT_MESSAGE_FORMAT, stream=None):
-        """ called from main after parsing arguments """
+        """
+        This function initializes the root logger and enables logging.
+        We set the level of the root logger to the one passed by calling logging.basicConfig.
+        We also set the level of every logger we created to 0 (logging.NOTSET), meaning that
+        the level of the root logger will be used to tell if messages should be logged.
+        Additionally, since our loggers use the NullHandler, they won't log anything themselves,
+        but due to having propagation enabled they will pass messages to the root logger,
+        which in turn will log to the stream set in this function.
+        Since the root logger is the one doing the work, when using JSON we set its formatter
+        so that every message logged is JSON-compatible.
+        """
         if self._is_enabled:
             raise ValueError('re-enabling logging. Not sure whether that is ok...')
  
@@ -90,26 +97,23 @@ class LogHelper:
         self._is_enabled = True
  
         self._use_json = use_json
-
-        sys.excepthook = self._log_except_hook
+        sys.excepthook = self._get_except_hook(sys.excepthook)
  
         # since there could be loggers already created we go through all of them
-        # and set their formatters to our custom Json formatter
-        # also set their levels so they respect what the main module wants logged
+        # and set their levels to 0 so they will use the root logger's level
         for name in self._all_names:
             logger = self.get_or_create_silent_logger(name)
+            self._set_logger_level(logger, logging.NOTSET)
  
-            if self._use_json:
-                self._make_json(logger)
-
-        # print the start of the logging message list
+        # add a JSON formatter to the root logger, which will be used by every logger
         if self._use_json:
+            _root_logger_wrapper.set_formatter(JsonFormatter())
             print('[')
  
     def end_logging(self):
         """
-        Must be called at the end of the main function
-        if the caller wants json-compatible ouptut
+        Must be called at the end of the main function if the caller wants
+        json-compatible output
         """
         if not self._is_enabled:
             return
@@ -124,60 +128,60 @@ class LogHelper:
             print(']')
         self._use_json = False
  
-    def _log_except_hook(self, exctype, value, traceback):
+    def _get_except_hook(self, old_hook):
         """
-        Global hook for exceptions so we can always end logging
+        Global hook for exceptions so we can always end logging.
+        We wrap any hook currently set to avoid overwriting global hooks set by oletools.
+        Note that this is only called by enable_logging, which in turn is called by
+        the main() function in oletools' scripts. When scripts are being imported this
+        code won't execute and won't affect global hooks.
         """
-        self.end_logging()
-        sys.__excepthook__(exctype, value, traceback)
+        def hook(exctype, value, traceback):
+            self.end_logging()
+            old_hook(exctype, value, traceback)
+
+        return hook
  
     def _get_or_create_logger(self, name, level, handler=None):
         """
-        If a logger doesn't exist, we create it and set the handler,
-        if it given. This avoids messing with existing loggers.
-        If we are using json then we also skip adding a handler,
-        since it would be removed afterwards
+        Get or create a new logger. This newly created logger will have the
+        handler and level that was passed, but if it already exists it's not changed.
+        We also wrap the logger in an adapter so we can easily extend its functionality.
         """
  
         # logging.getLogger creates a logger if it doesn't exist,
         # so we need to check before calling it
-        if handler and not self._use_json and not self._log_exists(name):
+        if handler and not self._log_exists(name):
             logger = logging.getLogger(name)
             logger.addHandler(handler)
+            self._set_logger_level(logger, level)
         else:
             logger = logging.getLogger(name)
  
-        self._set_logger_level(logger, level)
+        # Keep track of every logger we created so we can easily change
+        # their levels whenever needed
         self._all_names.add(name)
  
-        if self._use_json:
-            self._make_json(logger)
-
-        return logger
-
-    def _make_json(self, logger):
-        """
-        Replace handlers of every logger by a handler that uses the JSON formatter
-        """
-        for handler in logger.handlers:
-            logger.removeHandler(handler)
-        new_handler = logging.StreamHandler(sys.stdout)
-        new_handler.setFormatter(self._json_formatter)
-        logger.addHandler(new_handler)
+        adapted_logger = OletoolsLoggerAdapter(logger, None)
+        adapted_logger.set_json_enabled_function(lambda: self._use_json)
  
-        # Don't let it propagate to parent loggers, otherwise we might get
-        # duplicated messages from the root logger
-        logger.propagate = False
+        return adapted_logger
  
     @staticmethod
     def _set_logger_level(logger, level):
         """
-        If the logging is already initialized, we use the same level that
-        was set to the root logger. This prevents imported modules' loggers
-        from messing with the main module logging.
+        If the logging is already initialized, we set the level of our logger
+        to 0, meaning that it will reuse the level of the root logger.
+        That means that if the root logger level changes, we will keep using
+        its level and not logging unnecessarily.
         """
+
+        # if this log was wrapped, unwrap it to set the level
+        if isinstance(logger, OletoolsLoggerAdapter):
+            logger = logger.logger
+
         if _root_logger_wrapper.is_logging_initialized():
-            logger.setLevel(_root_logger_wrapper.get_root_logger_level())
+            logger.setLevel(logging.NOTSET)
         else:
             logger.setLevel(level)