pycolib/ansicolor.py
author Radek Brich <radek.brich@devl.cz>
Mon, 19 Aug 2013 14:18:26 +0200
changeset 9 8d5a1affbb9d
parent 5 055f7dfb3e4f
permissions -rw-r--r--
Update ansicolor: Support Python2. Add more documentation. Update handling of string color names.

# -*- coding: utf-8 -*-

# support Python 2.6+
from __future__ import print_function


#: Set disabled = True to disable all colors.
disabled = False


#: Supported colors as enum and string names.
(BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE, BOLD) = range(0,9)
names = ['black','red','green','yellow','blue','magenta','cyan','white','bold']


def highlight(enable, fg=None, bg=None):
    """Return ANSI escape code for selected color and font style.

    Set `enable` to true value to enable highlighting, false to disable/end highlighting.

    Set color and style using `fg` and `bg`. These parameters will accept
    string or variables defined in this module (BLACK, ...).

    Usage example::

        highlight(1) -- switch to bold, keep previous color
        highlight(1, 'bold red') -- bold red foreground
        highlight(1, BOLD | YELLOW, BLUE) -- yellow on blue, bold
        highlight(0) -- reset

    """
    global disabled
    if disabled:
        return ''

    if isinstance(fg, str):
        fg = sum(names.index(item.lower()) for item in fg.split())
    if isinstance(bg, str):
        bg = sum(names.index(item.lower()) for item in bg.split())

    if enable:
        code = '1'
        if fg is not None:
            code = ''
            if fg >= 8:
                fg -= 8
                code += '1;'
            code += str(30 + fg)
        if bg is not None:
            code += ';'
            if bg >= 8:
                bg -= 8
                code += '1;'
            code += str(40 + bg)
        return "\033[" + code + "m"
    else:
        return "\033[0m"


def hprint(*args, **kwargs):
    """Highlighted print.

    Usage is same as for print function.
    Whole output is colored as specified by `fg`, `bg` parameters.

    Usage example::

        hprint('Hello in red!', fg='bold red')

    """
    fg = kwargs.pop('fg', None)
    bg = kwargs.pop('bg', None)
    end = kwargs.pop('end', '\n')
    print(highlight(1, fg, bg), end='')
    print(*args, end='', **kwargs)
    print(highlight(0), end=end)