blob: fdb0fa6a8b4248eb549a84ac88e4234095d019b1 [file] [log] [blame]
# Copyright 2014 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
import contextlib
import logging
import os
from pylib.constants import host_paths
_COLORAMA_PATH = os.path.join(
host_paths.DIR_SOURCE_ROOT, 'third_party', 'colorama', 'src')
with host_paths.SysPath(_COLORAMA_PATH, position=0):
import colorama
BACK = colorama.Back
FORE = colorama.Fore
STYLE = colorama.Style
class _ColorFormatter(logging.Formatter):
# pylint does not see members added dynamically in the constructor.
# pylint: disable=no-member
color_map = {
logging.DEBUG: (FORE.CYAN),
logging.INFO: (), # Use default style.
logging.WARNING: (FORE.YELLOW),
logging.ERROR: (FORE.RED),
logging.CRITICAL: (BACK.RED),
}
def __init__(self, wrapped_formatter=None):
"""Wraps a |logging.Formatter| and adds color."""
super().__init__()
self._wrapped_formatter = wrapped_formatter or logging.Formatter()
#override
def format(self, record):
message = self._wrapped_formatter.format(record)
return self.Colorize(message, record.levelno)
def Colorize(self, message, log_level):
try:
return (''.join(self.color_map[log_level]) + message +
colorama.Style.RESET_ALL)
except KeyError:
return message
class ColorStreamHandler(logging.StreamHandler):
"""Handler that can be used to colorize logging output.
Example using a specific logger:
logger = logging.getLogger('my_logger')
logger.addHandler(ColorStreamHandler())
logger.info('message')
Example using the root logger:
ColorStreamHandler.MakeDefault()
logging.info('message')
"""
def __init__(self, force_color=False):
super().__init__()
self.force_color = force_color
self.setFormatter(logging.Formatter())
@property
def is_tty(self):
try:
isatty = getattr(self.stream, 'isatty')
except AttributeError:
return False
return isatty()
#override
def setFormatter(self, fmt):
if self.force_color or self.is_tty:
fmt = _ColorFormatter(fmt)
super().setFormatter(fmt)
@staticmethod
def MakeDefault(force_color=False):
"""
Replaces the default logging handlers with a coloring handler. To use
a colorizing handler at the same time as others, either register them
after this call, or add the ColorStreamHandler on the logger using
Logger.addHandler()
Args:
force_color: Set to True to bypass the tty check and always colorize.
"""
# If the existing handlers aren't removed, messages are duplicated
logging.getLogger().handlers = []
logging.getLogger().addHandler(ColorStreamHandler(force_color))
@contextlib.contextmanager
def OverrideColor(level, color):
"""Temporarily override the logging color for a specified level.
Args:
level: logging level whose color gets overridden.
color: tuple of formats to apply to log lines.
"""
prev_colors = {}
for handler in logging.getLogger().handlers:
if isinstance(handler.formatter, _ColorFormatter):
prev_colors[handler.formatter] = handler.formatter.color_map[level]
handler.formatter.color_map[level] = color
try:
yield
finally:
for formatter, prev_color in prev_colors.items():
formatter.color_map[level] = prev_color
@contextlib.contextmanager
def SuppressLogging(level=logging.ERROR):
"""Momentarilly suppress logging events from all loggers.
TODO(jbudorick): This is not thread safe. Log events from other threads might
also inadvertently disappear.
Example:
with logging_utils.SuppressLogging():
# all but CRITICAL logging messages are suppressed
logging.info('just doing some thing') # not shown
logging.critical('something really bad happened') # still shown
Args:
level: logging events with this or lower levels are suppressed.
"""
logging.disable(level)
yield
logging.disable(logging.NOTSET)