1 # Copyright (C) 2012 W. Trevor King <wking@drexel.edu>
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/>.
20 # Define ANSI escape sequences for colors
22 'black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white']
27 def standard_colors(use_color=None):
28 """Return a list of standard colors
30 >>> highlight,lowlight,good,bad = standard_colors()
31 >>> (highlight,lowlight,good,bad)
32 (None, 'blue', 'green', 'red')
42 highlight = lowlight = good = bad = None
43 return (highlight, lowlight, good, bad)
45 def _ansi_color_code(color):
46 r"""Return the appropriate ANSI escape sequence for `color`
48 >>> _ansi_color_code('blue')
50 >>> _ansi_color_code(None)
55 return '\033[3%dm' % (_COLORS.index(color))
57 def color_string(string, color=None):
58 r"""Wrap a string in ANSI escape sequences for coloring
60 >>> color_string('Hello world', 'red')
61 '\x1b[31mHello world\x1b[0m'
62 >>> color_string('Hello world', None)
65 It also works with non-unicode input:
67 >>> color_string('Hello world', 'red')
68 '\x1b[31mHello world\x1b[0m'
72 ret.append(_ansi_color_code(color))
75 ret.append(_ansi_color_code(None))
77 if isinstance(string, str): # i.e., not unicode
78 ret = [str(x) for x in ret]
82 def write_color(string, color=None, stream=None):
83 r"""Write a colored `string` to `stream`
85 If `stream` is `None`, it defaults to stdout.
87 >>> write_color('Hello world\n')
90 >>> from io import StringIO
91 >>> stream = StringIO()
92 >>> write_color('Hello world\n', 'red', stream)
94 '\x1b[31mHello world\n\x1b[0m'
98 stream.write(color_string(string=string, color=color))