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 class AnsiCmd(object):
  55     def __init__(self, forceAnsi):
  56         self.forceAnsi = forceAnsi
  57
  58     def cmdReset(self):
  59         ''' Returns the ansi cmd colour for a RESET '''
  60         if sys.stdout.isatty() or self.forceAnsi:
  61             return ESC + "[0m"
  62         else:
  63             return ""
  64
  65     def cmdColour(self, colour):
  66         '''
  67         Return the ansi cmd colour (i.e. escape sequence)
  68         for the ansi `colour` value
  69         '''
  70         if sys.stdout.isatty() or self.forceAnsi:
  71             return ESC + "[" + colour + "m"
  72         else:
  73             return ""
  74
  75     def cmdColourNamed(self, colour):
  76         ''' Return the ansi cmdColour for a given named `colour` '''
  77         try:
  78             return self.cmdColour(COLOURS_NAMED[colour])
  79         except KeyError:
  80             raise AnsiColourException('Unknown Colour %s' % (colour))
  81
  82     def cmdBold(self):
  83         if sys.stdout.isatty() or self.forceAnsi:
  84             return ESC + "[1m"
  85         else:
  86             return ""
  87
  88     def cmdUnderline(self):
  89         if sys.stdout.isatty() or self.forceAnsi:
  90             return ESC + "[4m"
  91         else:
  92             return ""
  93
  94 """These exist to maintain compatibility with users of version<=1.9.0"""
  95 def cmdReset():
  96     return AnsiCmd(False).cmdReset()
  97
  98 def cmdColour(colour):
  99     return AnsiCmd(False).cmdColour(colour)
 100
 101 def cmdColourNamed(colour):
 102     return AnsiCmd(False).cmdColourNamed(colour)