Ben Murdoch | 097c5b2 | 2016-05-18 11:27:45 +0100 | [diff] [blame^] | 1 | # Copyright 2014 The Chromium Authors. All rights reserved. |
| 2 | # Use of this source code is governed by a BSD-style license that can be |
| 3 | # found in the LICENSE file. |
| 4 | |
| 5 | import contextlib |
| 6 | import logging |
| 7 | import os |
| 8 | |
| 9 | from pylib.constants import host_paths |
| 10 | |
| 11 | _COLORAMA_PATH = os.path.join( |
| 12 | host_paths.DIR_SOURCE_ROOT, 'third_party', 'colorama', 'src') |
| 13 | |
| 14 | with host_paths.SysPath(_COLORAMA_PATH): |
| 15 | import colorama |
| 16 | |
| 17 | class ColorStreamHandler(logging.StreamHandler): |
| 18 | """Handler that can be used to colorize logging output. |
| 19 | |
| 20 | Example using a specific logger: |
| 21 | |
| 22 | logger = logging.getLogger('my_logger') |
| 23 | logger.addHandler(ColorStreamHandler()) |
| 24 | logger.info('message') |
| 25 | |
| 26 | Example using the root logger: |
| 27 | |
| 28 | ColorStreamHandler.MakeDefault() |
| 29 | logging.info('message') |
| 30 | |
| 31 | """ |
| 32 | # pylint does not see members added dynamically in the constructor. |
| 33 | # pylint: disable=no-member |
| 34 | color_map = { |
| 35 | logging.DEBUG: colorama.Fore.CYAN, |
| 36 | logging.WARNING: colorama.Fore.YELLOW, |
| 37 | logging.ERROR: colorama.Fore.RED, |
| 38 | logging.CRITICAL: colorama.Back.RED + colorama.Style.BRIGHT, |
| 39 | } |
| 40 | |
| 41 | def __init__(self, force_color=False): |
| 42 | super(ColorStreamHandler, self).__init__() |
| 43 | self.force_color = force_color |
| 44 | |
| 45 | @property |
| 46 | def is_tty(self): |
| 47 | isatty = getattr(self.stream, 'isatty', None) |
| 48 | return isatty and isatty() |
| 49 | |
| 50 | #override |
| 51 | def format(self, record): |
| 52 | message = logging.StreamHandler.format(self, record) |
| 53 | if self.force_color or self.is_tty: |
| 54 | return self.Colorize(message, record.levelno) |
| 55 | return message |
| 56 | |
| 57 | def Colorize(self, message, log_level): |
| 58 | try: |
| 59 | return self.color_map[log_level] + message + colorama.Style.RESET_ALL |
| 60 | except KeyError: |
| 61 | return message |
| 62 | |
| 63 | @staticmethod |
| 64 | def MakeDefault(force_color=False): |
| 65 | """ |
| 66 | Replaces the default logging handlers with a coloring handler. To use |
| 67 | a colorizing handler at the same time as others, either register them |
| 68 | after this call, or add the ColorStreamHandler on the logger using |
| 69 | Logger.addHandler() |
| 70 | |
| 71 | Args: |
| 72 | force_color: Set to True to bypass the tty check and always colorize. |
| 73 | """ |
| 74 | # If the existing handlers aren't removed, messages are duplicated |
| 75 | logging.getLogger().handlers = [] |
| 76 | logging.getLogger().addHandler(ColorStreamHandler(force_color)) |
| 77 | |
| 78 | |
| 79 | @contextlib.contextmanager |
| 80 | def SuppressLogging(level=logging.ERROR): |
| 81 | """Momentarilly suppress logging events from all loggers. |
| 82 | |
| 83 | TODO(jbudorick): This is not thread safe. Log events from other threads might |
| 84 | also inadvertently dissapear. |
| 85 | |
| 86 | Example: |
| 87 | |
| 88 | with logging_utils.SuppressLogging(): |
| 89 | # all but CRITICAL logging messages are suppressed |
| 90 | logging.info('just doing some thing') # not shown |
| 91 | logging.critical('something really bad happened') # still shown |
| 92 | |
| 93 | Args: |
| 94 | level: logging events with this or lower levels are suppressed. |
| 95 | """ |
| 96 | logging.disable(level) |
| 97 | yield |
| 98 | logging.disable(logging.NOTSET) |