From ebab1b6c69e3b5cd436bfaf31c957c3df1649c64 Mon Sep 17 00:00:00 2001
From: Thomas Hisch <t.hisch@gmail.com>
Date: Tue, 23 Jan 2018 00:57:56 +0100
Subject: [PATCH] live-logging: Colorize levelname

---
 _pytest/logging.py                | 60 ++++++++++++++++++++++++++++++-
 changelog/3142.feature            |  1 +
 testing/logging/test_formatter.py | 29 +++++++++++++++
 3 files changed, 89 insertions(+), 1 deletion(-)
 create mode 100644 changelog/3142.feature
 create mode 100644 testing/logging/test_formatter.py

diff --git a/_pytest/logging.py b/_pytest/logging.py
index f92b4c75e..095115cd9 100644
--- a/_pytest/logging.py
+++ b/_pytest/logging.py
@@ -2,8 +2,10 @@ from __future__ import absolute_import, division, print_function
 
 import logging
 from contextlib import closing, contextmanager
+import re
 import six
 
+from _pytest.config import create_terminal_writer
 import pytest
 import py
 
@@ -12,6 +14,58 @@ DEFAULT_LOG_FORMAT = '%(filename)-25s %(lineno)4d %(levelname)-8s %(message)s'
 DEFAULT_LOG_DATE_FORMAT = '%H:%M:%S'
 
 
+class ColoredLevelFormatter(logging.Formatter):
+    """
+    Colorize the %(levelname)..s part of the log format passed to __init__.
+    """
+
+    LOGLEVEL_COLOROPTS = {
+        logging.CRITICAL: {'red'},
+        logging.ERROR: {'red', 'bold'},
+        logging.WARNING: {'yellow'},
+        logging.WARN: {'yellow'},
+        logging.INFO: {'green'},
+        logging.DEBUG: {'purple'},
+        logging.NOTSET: set(),
+    }
+    LEVELNAME_FMT_REGEX = re.compile(r'%\(levelname\)([+-]?\d*s)')
+
+    def __init__(self, terminalwriter, *args, **kwargs):
+        super(ColoredLevelFormatter, self).__init__(
+            *args, **kwargs)
+        if six.PY2:
+            self._original_fmt = self._fmt
+        else:
+            self._original_fmt = self._style._fmt
+        self._level_to_fmt_mapping = {}
+
+        levelname_fmt_match = self.LEVELNAME_FMT_REGEX.search(self._fmt)
+        if not levelname_fmt_match:
+            return
+        levelname_fmt = levelname_fmt_match.group()
+
+        for level, color_opts in self.LOGLEVEL_COLOROPTS.items():
+            formatted_levelname = levelname_fmt % {
+                'levelname': logging.getLevelName(level)}
+
+            # add ANSI escape sequences around the formatted levelname
+            color_kwargs = {name: True for name in color_opts}
+            colorized_formatted_levelname = terminalwriter.markup(
+                formatted_levelname, **color_kwargs)
+            self._level_to_fmt_mapping[level] = self.LEVELNAME_FMT_REGEX.sub(
+                colorized_formatted_levelname,
+                self._fmt)
+
+    def format(self, record):
+        fmt = self._level_to_fmt_mapping.get(
+            record.levelno, self._original_fmt)
+        if six.PY2:
+            self._fmt = fmt
+        else:
+            self._style._fmt = fmt
+        return super(ColoredLevelFormatter, self).format(record)
+
+
 def get_option_ini(config, *names):
     for name in names:
         ret = config.getoption(name)  # 'default' arg won't work as expected
@@ -376,7 +430,11 @@ class LoggingPlugin(object):
             log_cli_handler = _LiveLoggingStreamHandler(terminal_reporter, capture_manager)
             log_cli_format = get_option_ini(self._config, 'log_cli_format', 'log_format')
             log_cli_date_format = get_option_ini(self._config, 'log_cli_date_format', 'log_date_format')
-            log_cli_formatter = logging.Formatter(log_cli_format, datefmt=log_cli_date_format)
+            if self._config.option.color != 'no' and ColoredLevelFormatter.LEVELNAME_FMT_REGEX.search(log_cli_format):
+                log_cli_formatter = ColoredLevelFormatter(create_terminal_writer(self._config),
+                                                          log_cli_format, datefmt=log_cli_date_format)
+            else:
+                log_cli_formatter = logging.Formatter(log_cli_format, datefmt=log_cli_date_format)
             log_cli_level = get_actual_log_level(self._config, 'log_cli_level', 'log_level')
             self.log_cli_handler = log_cli_handler
             self.live_logs_context = catching_logs(log_cli_handler, formatter=log_cli_formatter, level=log_cli_level)
diff --git a/changelog/3142.feature b/changelog/3142.feature
new file mode 100644
index 000000000..1461be514
--- /dev/null
+++ b/changelog/3142.feature
@@ -0,0 +1 @@
+Colorize the levelname column in the live-log output.
\ No newline at end of file
diff --git a/testing/logging/test_formatter.py b/testing/logging/test_formatter.py
new file mode 100644
index 000000000..10a921470
--- /dev/null
+++ b/testing/logging/test_formatter.py
@@ -0,0 +1,29 @@
+import logging
+
+import py.io
+from _pytest.logging import ColoredLevelFormatter
+
+
+def test_coloredlogformatter():
+    logfmt = '%(filename)-25s %(lineno)4d %(levelname)-8s %(message)s'
+
+    record = logging.LogRecord(
+        name='dummy', level=logging.INFO, pathname='dummypath', lineno=10,
+        msg='Test Message', args=(), exc_info=False)
+
+    class ColorConfig(object):
+        class option(object):
+            pass
+
+    tw = py.io.TerminalWriter()
+    tw.hasmarkup = True
+    formatter = ColoredLevelFormatter(tw, logfmt)
+    output = formatter.format(record)
+    assert output == ('dummypath                   10 '
+                      '\x1b[32mINFO    \x1b[0m Test Message')
+
+    tw.hasmarkup = False
+    formatter = ColoredLevelFormatter(tw, logfmt)
+    output = formatter.format(record)
+    assert output == ('dummypath                   10 '
+                      'INFO     Test Message')