1 # Copyright (C) 2012 W. Trevor King <wking@tremily.us>
3 # This file is part of pygrader.
5 # pygrader is free software: you can redistribute it and/or modify it under the
6 # terms of the GNU General Public License as published by the Free Software
7 # Foundation, either version 3 of the License, or (at your option) any later
10 # pygrader is distributed in the hope that it will be useful, but WITHOUT ANY
11 # WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
12 # A PARTICULAR PURPOSE. See the GNU General Public License for more details.
14 # You should have received a copy of the GNU General Public License along with
15 # pygrader. If not, see <http://www.gnu.org/licenses/>.
17 import logging as _logging
21 # Define ANSI escape sequences for colors
23 'black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white']
26 GOOD_DEBUG = _logging.DEBUG + 2
27 BAD_DEBUG = _logging.DEBUG + 4
30 def standard_colors(use_color=None):
31 """Return a list of standard colors
33 >>> highlight,lowlight,good,bad = standard_colors()
34 >>> (highlight,lowlight,good,bad)
35 (None, 'blue', 'green', 'red')
45 highlight = lowlight = good = bad = None
46 return (highlight, lowlight, good, bad)
48 def _ansi_color_code(color):
49 r"""Return the appropriate ANSI escape sequence for `color`
51 >>> _ansi_color_code('blue')
53 >>> _ansi_color_code(None)
58 return '\033[3%dm' % (_COLORS.index(color))
60 def color_string(string, color=None):
61 r"""Wrap a string in ANSI escape sequences for coloring
63 >>> color_string('Hello world', 'red')
64 '\x1b[31mHello world\x1b[0m'
65 >>> color_string('Hello world', None)
68 It also works with non-unicode input:
70 >>> color_string('Hello world', 'red')
71 '\x1b[31mHello world\x1b[0m'
75 ret.append(_ansi_color_code(color))
78 ret.append(_ansi_color_code(None))
80 if isinstance(string, str): # i.e., not unicode
81 ret = [str(x) for x in ret]
85 def write_color(string, color=None, stream=None):
86 r"""Write a colored `string` to `stream`
88 If `stream` is `None`, it defaults to stdout.
90 >>> write_color('Hello world\n')
93 >>> from io import StringIO
94 >>> stream = StringIO()
95 >>> write_color('Hello world\n', 'red', stream)
97 '\x1b[31mHello world\n\x1b[0m'
101 stream.write(color_string(string=string, color=color))
105 class ColoredFormatter (_logging.Formatter):
106 def __init__(self, *args, **kwargs):
107 super(ColoredFormatter, self).__init__(*args, **kwargs)
108 self.colored = None # `None` to use USE_COLOR; True/False to override
110 def format(self, record):
111 s = super(ColoredFormatter, self).format(record)
112 if self.colored or (self.colored is None and USE_COLOR):
113 highlight,lowlight,good,bad = standard_colors()
114 if record.levelno <= _logging.DEBUG:
116 elif record.levelno <= GOOD_DEBUG:
118 elif record.levelno <= BAD_DEBUG:
120 elif record.levelno <= _logging.INFO:
124 return color_string(string=s, color=color)