X-Git-Url: https://jfr.im/git/z_archive/twitter.git/blobdiff_plain/662f757601dc4f7db85c09211b9e81359106942f..a72824526aa3f87b5aa5cf5121f62f71aca42269:/twitter/ansi.py

diff --git a/twitter/ansi.py b/twitter/ansi.py
index 7b8f6db..1537000 100644
--- a/twitter/ansi.py
+++ b/twitter/ansi.py
@@ -1,6 +1,17 @@
 """
 Support for ANSI colours in command-line client.
 
+.. data:: ESC
+    ansi escape character
+
+.. data:: RESET
+    ansi reset colour (ansi value)
+
+.. data:: COLOURS_NAMED
+    dict of colour names mapped to their ansi value
+
+.. data:: COLOURS_MIDS
+    A list of ansi values for Mid Spectrum Colours
 """
 
 import itertools
@@ -9,26 +20,57 @@ import sys
 ESC = chr(0x1B)
 RESET = "0"
 
-COLOURS = [str(x) for x in range(31, 37)]
+COLOURS_NAMED = dict(list(zip(
+    ['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'],
+    [str(x) for x in range(30, 38)]
+)))
+COLOURS_MIDS = [
+    colour for name, colour in list(COLOURS_NAMED.items())
+    if name not in ('black', 'white')
+]
+
+class AnsiColourException(Exception):
+    ''' Exception while processing ansi colours '''
+    pass
 
 class ColourMap(object):
-    def __init__(self):
+    '''
+    Object that allows for mapping strings to ansi colour values.
+    '''
+    def __init__(self, colors=COLOURS_MIDS):
+        ''' uses the list of ansi `colors` values to initialize the map '''
         self._cmap = {}
-        self._colourIter = itertools.cycle(COLOURS)
-        
+        self._colourIter = itertools.cycle(colors)
+
     def colourFor(self, string):
-        if not self._cmap.has_key(string):
-            self._cmap[string] = self._colourIter.next()
+        '''
+        Returns an ansi colour value given a `string`.
+        The same ansi colour value is always returned for the same string
+        '''
+        if string not in self._cmap:
+            self._cmap[string] = next(self._colourIter)
         return self._cmap[string]
 
 def cmdReset():
+    ''' Returns the ansi cmd colour for a RESET '''
     if sys.stdout.isatty():
         return ESC + "[0m"
     else:
         return ""
 
 def cmdColour(colour):
+    '''
+    Return the ansi cmd colour (i.e. escape sequence)
+    for the ansi `colour` value
+    '''
     if sys.stdout.isatty():
         return ESC + "[" + colour + "m"
     else:
         return ""
+
+def cmdColourNamed(colour):
+    ''' Return the ansi cmdColour for a given named `colour` '''
+    try:
+        return cmdColour(COLOURS_NAMED[colour])
+    except KeyError:
+        raise AnsiColourException('Unknown Colour %s' %(colour))