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
  19 ESC = chr(0x1B)
  20 RESET = "0"
  21
  22 COLOURS_NAMED = dict(zip(
  23     ['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'],
  24     [str(x) for x in range(30, 38)]
  25 ))
  26 COLOURS_MIDS = [
  27     colour for name, colour in COLOURS_NAMED.items()
  28     if name not in ('black', 'white')
  29 ]
  30
  31 class AnsiColourException(Exception):
  32     ''' Exception while processing ansi colours '''
  33     pass
  34
  35 class ColourMap(object):
  36     '''
  37     Object that allows for mapping strings to ansi colour values.
  38     '''
  39     def __init__(self, colors=COLOURS_MIDS):
  40         ''' uses the list of ansi `colors` values to initialize the map '''
  41         self._cmap = {}
  42         self._colourIter = itertools.cycle(colors)
  43
  44     def colourFor(self, string):
  45         '''
  46         Returns an ansi colour value given a `string`.
  47         The same ansi colour value is always returned for the same string
  48         '''
  49         if not self._cmap.has_key(string):
  50             self._cmap[string] = self._colourIter.next()
  51         return self._cmap[string]
  52
  53 def cmdColour(colour):
  54     '''
  55     Return the ansi cmd colour (i.e. escape sequence)
  56     for the ansi `colour` value
  57     '''
  58     return ESC + "[" + colour + "m"
  59
  60 def cmdReset():
  61     ''' Returns the ansi cmd colour for a RESET '''
  62     return cmdColour(RESET)
  63
  64 def cmdColourNamed(colour):
  65     ''' Return the ansi cmdColour for a given named `colour` '''
  66     try:
  67         return cmdColour(COLOURS_NAMED[colour])
  68     except KeyError:
  69         raise AnsiColourException('Unknown Colour %s' %(colour))