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 logging as _logging
  18 import sys as _sys
  19
  20
  21 # Define ANSI escape sequences for colors
  22 _COLORS = [
  23     'black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white']
  24
  25 USE_COLOR = True
  26 GOOD_DEBUG = _logging.DEBUG + 2
  27 BAD_DEBUG = _logging.DEBUG + 4
  28
  29
  30 def standard_colors(use_color=None):
  31     """Return a list of standard colors
  32
  33     >>> highlight,lowlight,good,bad = standard_colors()
  34     >>> (highlight,lowlight,good,bad)
  35     (None, 'blue', 'green', 'red')
  36     """
  37     if use_color is None:
  38         use_color = USE_COLOR
  39     if use_color:
  40         highlight = None
  41         lowlight = 'blue'
  42         good = 'green'
  43         bad = 'red'
  44     else:
  45         highlight = lowlight = good = bad = None
  46     return (highlight, lowlight, good, bad)
  47
  48 def _ansi_color_code(color):
  49     r"""Return the appropriate ANSI escape sequence for `color`
  50
  51     >>> _ansi_color_code('blue')
  52     '\x1b[34m'
  53     >>> _ansi_color_code(None)
  54     '\x1b[0m'
  55     """
  56     if color is None:
  57         return '\033[0m'
  58     return '\033[3%dm' % (_COLORS.index(color))
  59
  60 def color_string(string, color=None):
  61     r"""Wrap a string in ANSI escape sequences for coloring
  62
  63     >>> color_string('Hello world', 'red')
  64     '\x1b[31mHello world\x1b[0m'
  65     >>> color_string('Hello world', None)
  66     'Hello world'
  67
  68     It also works with non-unicode input:
  69
  70     >>> color_string('Hello world', 'red')
  71     '\x1b[31mHello world\x1b[0m'
  72     """
  73     ret = []
  74     if color:
  75         ret.append(_ansi_color_code(color))
  76     ret.append(string)
  77     if color:
  78         ret.append(_ansi_color_code(None))
  79     sep = ''
  80     if isinstance(string, str):  # i.e., not unicode
  81         ret = [str(x) for x in ret]
  82         sep = ''
  83     return sep.join(ret)
  84
  85 def write_color(string, color=None, stream=None):
  86     r"""Write a colored `string` to `stream`
  87
  88     If `stream` is `None`, it defaults to stdout.
  89
  90     >>> write_color('Hello world\n')
  91     Hello world
  92
  93     >>> from io import StringIO
  94     >>> stream = StringIO()
  95     >>> write_color('Hello world\n', 'red', stream)
  96     >>> stream.getvalue()
  97     '\x1b[31mHello world\n\x1b[0m'
  98     """
  99     if stream is None:
 100         stream = _sys.stdout
 101     stream.write(color_string(string=string, color=color))
 102     stream.flush()
 103
 104
 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
 109
 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:
 115                 color = lowlight
 116             elif record.levelno <= GOOD_DEBUG:
 117                 color = good
 118             elif record.levelno <= BAD_DEBUG:
 119                 color = bad
 120             elif record.levelno <= _logging.INFO:
 121                 color = highlight
 122             else:
 123                 color = bad
 124             return color_string(string=s, color=color)
 125         return s