Create a prefix logger adapter class to add prefix to logs (#849)

2 files changed, 61 insertions, 1 deletions

diff --git a/mobly/logger.py b/mobly/logger.py
index 27f747c..365f5c5 100644
--- a/mobly/logger.py
+++ b/mobly/logger.py
@@ -17,6 +17,7 @@ import logging
 import os
 import re
 import sys
+from typing import Any, MutableMapping, Tuple
 
 from mobly import records
 from mobly import utils
@@ -373,3 +374,49 @@ def normalize_log_line_timestamp(log_line_timestamp):
     special characters.
   """
   return sanitize_filename(log_line_timestamp)
+
+
+class PrefixLoggerAdapter(logging.LoggerAdapter):
+  """A wrapper that adds a prefix to each log line.
+
+  This logger adapter class is like a decorator to Logger. It takes one
+  Logger-like object and returns a new Logger-like object. The new Logger-like
+  object will print logs with a custom prefix added. Creating new Logger-like
+  objects doesn't modify the behavior of the old Logger-like object.
+
+  Chaining multiple logger adapters is also supported. The multiple adapters
+  will take effect in the order in which they are chained, i.e. the log will be
+  '<prefix1> <prefix2> <prefix3> <message>' if we chain 3 PrefixLoggerAdapters.
+
+  Example Usage:
+
+  .. code-block:: python
+
+    logger = PrefixLoggerAdapter(logging.getLogger(), {
+      'log_prefix': <custom prefix>
+    })
+
+  Then each log line added by the logger will have a prefix:
+  '<custom prefix> <message>'.
+  """
+
+  _KWARGS_TYPE = MutableMapping[str, Any]
+  _PROCESS_RETURN_TYPE = Tuple[str, _KWARGS_TYPE]
+
+  # The key of log_preifx item in the dict self.extra
+  EXTRA_KEY_LOG_PREFIX: str = 'log_prefix'
+
+  extra: _KWARGS_TYPE
+
+  def process(self, msg: str, kwargs: _KWARGS_TYPE) -> _PROCESS_RETURN_TYPE:
+    """Processes the logging call to insert contextual information.
+
+    Args:
+      msg: the logging message
+      kwargs: keyword arguments passed in to a logging call
+
+    Returns:
+      the message and kwargs modified.
+    """
+    new_msg = f'{self.extra[PrefixLoggerAdapter.EXTRA_KEY_LOG_PREFIX]} {msg}'
+    return (new_msg, kwargs)
diff --git a/tests/mobly/logger_test.py b/tests/mobly/logger_test.py
index 1d568c5..8a1e778 100755
--- a/tests/mobly/logger_test.py
+++ b/tests/mobly/logger_test.py
@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import os
 import shutil
 import tempfile
 import unittest
-import logging
 from unittest import mock
 
 from mobly import logger
@@ -210,6 +210,19 @@ class LoggerTest(unittest.TestCase):
     expected_filename = 'logcat.txt_'
     self.assertEqual(logger.sanitize_filename(fake_filename), expected_filename)
 
+  def test_prefix_logger_adapter_prefix_log_lines(self):
+    extra = {
+        logger.PrefixLoggerAdapter.EXTRA_KEY_LOG_PREFIX: '[MOCK_PREFIX]',
+    }
+    adapted_logger = logger.PrefixLoggerAdapter(mock.Mock(), extra)
+
+    kwargs = mock.Mock()
+    processed_log, processed_kwargs = adapted_logger.process('mock log line',
+                                                             kwargs=kwargs)
+
+    self.assertEqual(processed_log, '[MOCK_PREFIX] mock log line')
+    self.assertIs(processed_kwargs, kwargs)
+
 
 if __name__ == "__main__":
   unittest.main()