Saturday, October 13, 2012

Color text output from python

At the moment I run a minimization routine build in Cython/Python that outputs a lot of text to the terminal. Not all text is equally interesting and I was wondering if it is difficult to output some of the more important messages in a different color. It turns out that this is very easy. On stackoverflow Dietrich Epp posted a nice example which I have changed a little to select the color with a single character string:

def hilite(string, status, bold):
    attr = []
    if(sys.stdout.isatty()):
        if status=='g':
            # green
            attr.append('32')
        elif status=='r':
            # red
            attr.append('31')
        elif status=='y':
            # yellow
            attr.append('33')    
        elif status=='b':
            # blue
            attr.append('34')
        elif status=='m':
            # magenta
            attr.append('35')
        if bold:
            attr.append('1')

        return '\x1b[%sm%s\x1b[0m' % (';'.join(attr), string)
    else:
        return(string)

More info on ANSI escape codes is given on Wikipedia. You can also do things like underline. So far I have only tested this on Linux so I do not know if it also works on Windows.

Update:
This does not work on Windows (windows terminals do not work with ANSI codes). Luckily there is a python package you can use to make it work on windows after all: colorama. Simply call colorama.init()
(after import colorama) and you can use the above code again. However, I noticed that if you use ipython this garbles your prompt (and other things). To work around this I call colorama.init() in the functions that use the ANSI codes in the output. Then when the script is finished you have to call colorama.deinit() and ipython will behave normally.

No comments:

Post a Comment