twitter/ansi.py

   1 """
   2 Support for ANSI colours in command-line client.
   3
   4 .. data:: ESC
   5     ansi escape character
   6
   7 .. data:: RESET
   8     ansi reset colour (ansi value)
   9
  10 .. data:: COLOURS_NAMED
  11     dict of colour names mapped to their ansi value
  12
  13 .. data:: COLOURS_MIDS
  14     A list of ansi values for Mid Spectrum Colours
  15 """
  16
  17 import itertools
  18 import sys
  19
  20 ESC = chr(0x1B)
  21 RESET = "0"
  22
  23 COLOURS_NAMED = dict(list(zip(
  24     ['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'],
  25     [str(x) for x in range(30, 38)]
  26 )))
  27 COLOURS_MIDS = [
  28     colour for name, colour in list(COLOURS_NAMED.items())
  29     if name not in ('black', 'white')
  30 ]
  31
  32 class AnsiColourException(Exception):
  33     ''' Exception while processing ansi colours '''
  34     pass
  35
  36 class ColourMap(object):
  37     '''
  38     Object that allows for mapping strings to ansi colour values.
  39     '''
  40     def __init__(self, colors=COLOURS_MIDS):
  41         ''' uses the list of ansi `colors` values to initialize the map '''
  42         self._cmap = {}
  43         self._colourIter = itertools.cycle(colors)
  44
  45     def colourFor(self, string):
  46         '''
  47         Returns an ansi colour value given a `string`.
  48         The same ansi colour value is always returned for the same string
  49         '''
  50         if string not in self._cmap:
  51             self._cmap[string] = next(self._colourIter)
  52         return self._cmap[string]
  53
  54 def cmdReset():
  55     ''' Returns the ansi cmd colour for a RESET '''
  56     if sys.stdout.isatty():
  57         return ESC + "[0m"
  58     else:
  59         return ""
  60
  61 def cmdColour(colour):
  62     '''
  63     Return the ansi cmd colour (i.e. escape sequence)
  64     for the ansi `colour` value
  65     '''
  66     if sys.stdout.isatty():
  67         return ESC + "[" + colour + "m"
  68     else:
  69         return ""
  70
  71 def cmdColourNamed(colour):
  72     ''' Return the ansi cmdColour for a given named `colour` '''
  73     try:
  74         return cmdColour(COLOURS_NAMED[colour])
  75     except KeyError:
  76         raise AnsiColourException('Unknown Colour %s' %(colour))