1"""Nicer log formatting with colours.
2
3Code copied from Tornado, Apache licensed.
4"""
5# Copyright 2012 Facebook
6#
7# Licensed under the Apache License, Version 2.0 (the "License"); you may
8# not use this file except in compliance with the License. You may obtain
9# a copy of the License at
10#
11# http://www.apache.org/licenses/LICENSE-2.0
12#
13# Unless required by applicable law or agreed to in writing, software
14# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
15# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
16# License for the specific language governing permissions and limitations
17# under the License.
18
19import logging
20import sys
21
22try:
23    import curses
24except ImportError:
25    curses = None
26
27
28def _stderr_supports_color():
29    color = False
30    if curses and hasattr(sys.stderr, 'isatty') and sys.stderr.isatty():
31        try:
32            curses.setupterm()
33            if curses.tigetnum("colors") > 0:
34                color = True
35        except Exception:
36            pass
37    return color
38
39
40class LogFormatter(logging.Formatter):
41    """Log formatter with colour support
42    """
43    DEFAULT_COLORS = {
44        logging.INFO: 2,  # Green
45        logging.WARNING: 3,  # Yellow
46        logging.ERROR: 1,  # Red
47        logging.CRITICAL: 1,
48    }
49
50    def __init__(self, color=True, datefmt=None):
51        r"""
52        :arg bool color: Enables color support.
53        :arg string fmt: Log message format.
54        It will be applied to the attributes dict of log records. The
55        text between ``%(color)s`` and ``%(end_color)s`` will be colored
56        depending on the level if color support is on.
57        :arg dict colors: color mappings from logging level to terminal color
58        code
59        :arg string datefmt: Datetime format.
60        Used for formatting ``(asctime)`` placeholder in ``prefix_fmt``.
61        .. versionchanged:: 3.2
62        Added ``fmt`` and ``datefmt`` arguments.
63        """
64        logging.Formatter.__init__(self, datefmt=datefmt)
65        self._colors = {}
66        if color and _stderr_supports_color():
67            # The curses module has some str/bytes confusion in
68            # python3. Until version 3.2.3, most methods return
69            # bytes, but only accept strings. In addition, we want to
70            # output these strings with the logging module, which
71            # works with unicode strings. The explicit calls to
72            # unicode() below are harmless in python2 but will do the
73            # right conversion in python 3.
74            fg_color = (curses.tigetstr("setaf") or
75                        curses.tigetstr("setf") or "")
76            if (3, 0) < sys.version_info < (3, 2, 3):
77                fg_color = str(fg_color, "ascii")
78
79            for levelno, code in self.DEFAULT_COLORS.items():
80                self._colors[levelno] = str(
81                    curses.tparm(fg_color, code), "ascii")
82            self._normal = str(curses.tigetstr("sgr0"), "ascii")
83
84            scr = curses.initscr()
85            self.termwidth = scr.getmaxyx()[1]
86            curses.endwin()
87        else:
88            self._normal = ''
89            # Default width is usually 80, but too wide is
90            # worse than too narrow
91            self.termwidth = 70
92
93    def formatMessage(self, record):
94        mlen = len(record.message)
95        right_text = '{initial}-{name}'.format(initial=record.levelname[0],
96                                               name=record.name)
97        if mlen + len(right_text) < self.termwidth:
98            space = ' ' * (self.termwidth - (mlen + len(right_text)))
99        else:
100            space = '  '
101
102        if record.levelno in self._colors:
103            start_color = self._colors[record.levelno]
104            end_color = self._normal
105        else:
106            start_color = end_color = ''
107
108        return record.message + space + start_color + right_text + end_color
109
110
111def enable_colourful_output(level=logging.INFO):
112    handler = logging.StreamHandler()
113    handler.setFormatter(LogFormatter())
114    logging.root.addHandler(handler)
115    logging.root.setLevel(level)
116