1# Copyright 2014 The Chromium Authors 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5import contextlib 6import logging 7import os 8 9from pylib.constants import host_paths 10 11_COLORAMA_PATH = os.path.join( 12 host_paths.DIR_SOURCE_ROOT, 'third_party', 'colorama', 'src') 13 14with host_paths.SysPath(_COLORAMA_PATH, position=0): 15 import colorama 16 17BACK = colorama.Back 18FORE = colorama.Fore 19STYLE = colorama.Style 20 21 22class _ColorFormatter(logging.Formatter): 23 # pylint does not see members added dynamically in the constructor. 24 # pylint: disable=no-member 25 color_map = { 26 logging.DEBUG: (FORE.CYAN), 27 logging.INFO: (), # Use default style. 28 logging.WARNING: (FORE.YELLOW), 29 logging.ERROR: (FORE.RED), 30 logging.CRITICAL: (BACK.RED), 31 } 32 33 def __init__(self, wrapped_formatter=None): 34 """Wraps a |logging.Formatter| and adds color.""" 35 super().__init__() 36 self._wrapped_formatter = wrapped_formatter or logging.Formatter() 37 38 #override 39 def format(self, record): 40 message = self._wrapped_formatter.format(record) 41 return self.Colorize(message, record.levelno) 42 43 def Colorize(self, message, log_level): 44 try: 45 return (''.join(self.color_map[log_level]) + message + 46 colorama.Style.RESET_ALL) 47 except KeyError: 48 return message 49 50 51class ColorStreamHandler(logging.StreamHandler): 52 """Handler that can be used to colorize logging output. 53 54 Example using a specific logger: 55 56 logger = logging.getLogger('my_logger') 57 logger.addHandler(ColorStreamHandler()) 58 logger.info('message') 59 60 Example using the root logger: 61 62 ColorStreamHandler.MakeDefault() 63 logging.info('message') 64 65 """ 66 def __init__(self, force_color=False): 67 super().__init__() 68 self.force_color = force_color 69 self.setFormatter(logging.Formatter()) 70 71 @property 72 def is_tty(self): 73 try: 74 isatty = getattr(self.stream, 'isatty') 75 except AttributeError: 76 return False 77 return isatty() 78 79 #override 80 def setFormatter(self, fmt): 81 if self.force_color or self.is_tty: 82 fmt = _ColorFormatter(fmt) 83 super().setFormatter(fmt) 84 85 @staticmethod 86 def MakeDefault(force_color=False): 87 """ 88 Replaces the default logging handlers with a coloring handler. To use 89 a colorizing handler at the same time as others, either register them 90 after this call, or add the ColorStreamHandler on the logger using 91 Logger.addHandler() 92 93 Args: 94 force_color: Set to True to bypass the tty check and always colorize. 95 """ 96 # If the existing handlers aren't removed, messages are duplicated 97 logging.getLogger().handlers = [] 98 logging.getLogger().addHandler(ColorStreamHandler(force_color)) 99 100 101@contextlib.contextmanager 102def OverrideColor(level, color): 103 """Temporarily override the logging color for a specified level. 104 105 Args: 106 level: logging level whose color gets overridden. 107 color: tuple of formats to apply to log lines. 108 """ 109 prev_colors = {} 110 for handler in logging.getLogger().handlers: 111 if isinstance(handler.formatter, _ColorFormatter): 112 prev_colors[handler.formatter] = handler.formatter.color_map[level] 113 handler.formatter.color_map[level] = color 114 try: 115 yield 116 finally: 117 for formatter, prev_color in prev_colors.items(): 118 formatter.color_map[level] = prev_color 119 120 121@contextlib.contextmanager 122def SuppressLogging(level=logging.ERROR): 123 """Momentarilly suppress logging events from all loggers. 124 125 TODO(jbudorick): This is not thread safe. Log events from other threads might 126 also inadvertently disappear. 127 128 Example: 129 130 with logging_utils.SuppressLogging(): 131 # all but CRITICAL logging messages are suppressed 132 logging.info('just doing some thing') # not shown 133 logging.critical('something really bad happened') # still shown 134 135 Args: 136 level: logging events with this or lower levels are suppressed. 137 """ 138 logging.disable(level) 139 yield 140 logging.disable(logging.NOTSET) 141