pygrader/color.py

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