pygrader/color.py

   1 # Copyright (C) 2012 W. Trevor King <wking@drexel.edu>
   2 #
   3 # This file is part of pygrader.
   4 #
   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
   8 # version.
   9 #
  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.
  13 #
  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/>.
  16
  17 import sys as _sys
  18
  19
  20 # Define ANSI escape sequences for colors
  21 _COLORS = [
  22     'black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white']
  23
  24 USE_COLOR = True
  25
  26
  27 def standard_colors(use_color=None):
  28     """Return a list of standard colors
  29
  30     >>> highlight,lowlight,good,bad = standard_colors()
  31     >>> (highlight,lowlight,good,bad)
  32     (None, 'blue', 'green', 'red')
  33     """
  34     if use_color is None:
  35         use_color = USE_COLOR
  36     if use_color:
  37         highlight = None
  38         lowlight = 'blue'
  39         good = 'green'
  40         bad = 'red'
  41     else:
  42         highlight = lowlight = good = bad = None
  43     return (highlight, lowlight, good, bad)
  44
  45 def _ansi_color_code(color):
  46     r"""Return the appropriate ANSI escape sequence for `color`
  47
  48     >>> _ansi_color_code('blue')
  49     '\x1b[34m'
  50     >>> _ansi_color_code(None)
  51     '\x1b[0m'
  52     """
  53     if color is None:
  54         return '\033[0m'
  55     return '\033[3%dm' % (_COLORS.index(color))
  56
  57 def color_string(string, color=None):
  58     r"""Wrap a string in ANSI escape sequences for coloring
  59
  60     >>> color_string('Hello world', 'red')
  61     '\x1b[31mHello world\x1b[0m'
  62     >>> color_string('Hello world', None)
  63     'Hello world'
  64
  65     It also works with non-unicode input:
  66
  67     >>> color_string('Hello world', 'red')
  68     '\x1b[31mHello world\x1b[0m'
  69     """
  70     ret = []
  71     if color:
  72         ret.append(_ansi_color_code(color))
  73     ret.append(string)
  74     if color:
  75         ret.append(_ansi_color_code(None))
  76     sep = ''
  77     if isinstance(string, str):  # i.e., not unicode
  78         ret = [str(x) for x in ret]
  79         sep = ''
  80     return sep.join(ret)
  81
  82 def write_color(string, color=None, stream=None):
  83     r"""Write a colored `string` to `stream`
  84
  85     If `stream` is `None`, it defaults to stdout.
  86
  87     >>> write_color('Hello world\n')
  88     Hello world
  89
  90     >>> from io import StringIO
  91     >>> stream = StringIO()
  92     >>> write_color('Hello world\n', 'red', stream)
  93     >>> stream.getvalue()
  94     '\x1b[31mHello world\n\x1b[0m'
  95     """
  96     if stream is None:
  97         stream = _sys.stdout
  98     stream.write(color_string(string=string, color=color))
  99     stream.flush()