"colorize" - The tiny module which provides ANSI color formatting

colorize.py

"""
This tiny module provides ANSI color formatting.

    ANSI escape code - Wikipedia, the free encyclopedia 
    http://en.wikipedia.org/wiki/ANSI_escape_code

>>> colorize("")
''
>>> colorize("black and white")
'black and white'
>>> colorize("black", COLOR_BLACK)
'\\x1b[30mblack\\x1b[0m'
>>> colorize("red and black", COLOR_RED, COLOR_BLACK)
'\\x1b[31;40mred and black\\x1b[0m'

"""

CSI="\x1B["

SET_FG_COLOR = 30
SET_BG_COLOR = 40

COLOR_BLACK   = 0
COLOR_RED     = 1
COLOR_GREEN   = 2
COLOR_YELLOW  = 3
COLOR_BLUE    = 4
COLOR_MAGENTA = 5
COLOR_CYAN    = 6
COLOR_WHITE   = 7


def colorize(string, fgcolor=None, bgcolor=None):
    sgr = []

    if fgcolor is not None:
        sgr.append(SET_FG_COLOR + fgcolor)
    if bgcolor is not None:
        sgr.append(SET_BG_COLOR + bgcolor)

    if not sgr or not string:
        return string

    sgr = [str(x) for x in sgr]
    fmt = [CSI, ';'.join(sgr), 'm', string, CSI, '0m']

    return ''.join(fmt)


if __name__ == "__main__":
    import doctest
    doctest.testmod()