colout/colout.py

#!/usr/bin/env python
#encoding: utf-8

# Color Up Arbitrary Command Ouput 
# Licensed under the GPL version 3
# 2012 (c) nojhan <nojhan@nojhan.net>

import re
import random

###########
# Library #
###########

# Available styles
styles = {
    "normal":0, "bold":1, "faint":2, "italic":3, "underline":4,
    "blink":5, "rapid_blink":6,
    "reverse":7, "conceal":8
}

# Available color names in 8-colors mode
colors = {
    "black":0, "red":1, "green":2, "yellow":3, "blue":4,
    "magenta":5, "cyan":6, "white":7
}

rainbow = [ "red", "yellow", "green", "cyan", "blue", "magenta" ]
colormap = rainbow  # default colormap to rainbow
colormap_idx = 0

# Escaped end markers for given color modes
endmarks = {8:";", 256:";38;5;"}

def colorin( text, color = "red", style = "normal" ):
    """
    Return the given text, surrounded by the given color ASCII markers.

    If the given color is a name that exists in available colors,
    a 8-colors mode is assumed, else, a 256-colors mode.

    The given style must exists in the available styles.

    >>> colorin("Fetchez la vache", "red", "bold")
    '\x1b[1;31mFetchez la vache\x1b[0m'
    >>> colout.colorin("Faites chier la vache", 41, "normal")
    '\x1b[0;38;5;41mFaites chier la vache\x1b[0m'
    """
    global colormap_idx

    # Special characters.
    start = "\033["
    stop = "\033[0m"

    # Convert the style code
    if style == "random" or style == "Random":
        style = random.choice(list(styles.keys()))
    else:
        assert( style in styles)

    style_code = str(styles[style])

    if color == "random":
        mode = 8
        color_code = random.choice(list(colors.values()))
        color_code = str( 30 + color_code )

    elif color == "Random":
        mode = 256
        color_nb = random.randint(0,255)
        color_code = str( color_nb )

    elif color == "rainbow":
        mode = 8
        color = colormap[colormap_idx]
        color_code = str( 30 + colors[color] )

        if colormap_idx < len(colormap)-1:
            colormap_idx += 1
        else:
            colormap_idx = 0

    elif color == "colormap":
        color = colormap[colormap_idx]
        if color in colors:
            mode = 8
            color_code = str( 30 + colors[color] )
        else:
            mode = 256
            color_nb = int( color )
            assert( 0 <= color_nb <= 255 )
            color_code = str( color_nb )

        if colormap_idx < len(colormap)-1:
            colormap_idx += 1
        else:
            colormap_idx = 0

    # 8 colors modes
    elif color in colors:
        mode = 8
        color_code = str( 30 + colors[color] )

    # 256 colors mode
    else:
        mode = 256
        color_nb = int( color )
        assert( 0 <= color_nb <= 255 )
        color_code = str( color_nb )

    return start + style_code + endmarks[mode] + color_code + "m" + text + stop


def colorout( text, match, prev_end, color = "red", style = "normal", group=0 ):
    """
    Build the text from the previous re.match to the current one,
    coloring up the matching characters.
    """
    start = match.start(group)
    colored_text = text[prev_end:start]
    end = match.end(group)

    colored_text += colorin(text[start:end], color, style )
    return colored_text,end


def colorup( text, pattern, color = "red", style = "normal", on_groups=False):
    """
    Color up every characters that match the given regexp patterns.
    If groups are specified, only color up them and not the whole pattern.

    Colors and styles may be specified as a list of comma-separated values,
    in which case the different matching groups may be formatted differently.
    If there is less colors/styles than groups, the last format is used
    for the additional groups.

    >>> colorup("Fetchez la vache", "vache", "red", "bold")
    'Fetchez la \x1b[1;31mvache\x1b[0m'
    >>> colorup("Faites chier la vache", "[Fv]a", "red", "bold")
    '\x1b[1;31mFa\x1b[0mites chier la \x1b[1;31mva\x1b[0mche'
    >>> colorup("Faites Chier la Vache", "[A-Z](\S+)\s", "red", "bold")
    'F\x1b[1;31maites\x1b[0m C\x1b[1;31mhier\x1b[0m la Vache'
    >>> colorup("Faites Chier la Vache", "([A-Z])(\S+)\s", "red,green", "bold")
    '\x1b[1;31mF\x1b[0m\x1b[1;32maites\x1b[0m \x1b[1;31mC\x1b[0m\x1b[1;32mhier\x1b[0m la Vache'
    >>> colorup("Faites Chier la Vache", "([A-Z])(\S+)\s", "green")
    '\x1b[0;32mF\x1b[0m\x1b[0;32maites\x1b[0m \x1b[0;32mC\x1b[0m\x1b[0;32mhier\x1b[0m la Vache'
    >>> colorup("Faites Chier la Vache", "([A-Z])(\S+)\s", "blue", "bold,italic")
    '\x1b[1;34mF\x1b[0m\x1b[3;34maites\x1b[0m \x1b[1;34mC\x1b[0m\x1b[3;34mhier\x1b[0m la Vache'
    """
    global colormap_idx
    regex = re.compile(pattern)#, re.IGNORECASE)

    # Prepare the colored text.
    colored_text = ""
    end = 0
    for match in regex.finditer(text):

        # If no groups are specified
        if not match.groups():
            # Color the previous partial line,
            partial,end = colorout( text, match, end, color, style )
            # add it to the final text.
            colored_text += partial

        else:
            nb_groups = len(match.groups())

            # Build a list of colors that match the number of grouped,
            # if there is not enough colors, duplicate the last one.
            colors_l = color.split(",")
            group_colors = colors_l + [colors_l[-1]] * (nb_groups - len(colors_l))

            # Same for styles
            styles_l = style.split(",")
            group_styles = styles_l + [styles_l[-1]] * (nb_groups - len(styles_l))

            # If we want to iterate colormaps on groups instead of patterns
            if on_groups:
                # Reset the counter at the beginning of each match
                colormap_idx = 0

            # For each group index.
            # Note that match.groups returns a tuple (thus being indexed in [0,n[),
            # but that match.start(0) refers to the whole match, the groups being indexed in [1,n].
            # Thus, we need to range in [1,n+1[.
            for group in range(1,nb_groups+1):
                partial,end = colorout( text, match, end, group_colors[group-1], group_styles[group-1], group )
                colored_text += partial

    # Append the remaining part of the text, if any.
    colored_text += text[end:]

    return colored_text


def colorgen( stream, pattern, color = "red", style = "normal", on_groups=False):
    """
    A generator that colors the items given in an iterable input.

    >>> import math
    >>> list(colorgen([str(i) for i in [math.pi,math.e]],"1","red"))
    ['3.\x1b[0;31m1\x1b[0m4\x1b[0;31m1\x1b[0m59265359',
     '2.7\x1b[0;31m1\x1b[0m828\x1b[0;31m1\x1b[0m82846']
    """
    while True:
        try:
            item = stream.readline()
        except KeyboardInterrupt:
            break
        if not item:
            break
        yield colorup( item, pattern, color, style, on_groups)


######################
# Command line tools #
######################

def __args_dirty__(argv,usage=""):
    """
    Roughly extract options from the command line arguments.
    To be used only when argparse is not available.

    Returns a tuple of (pattern,color,style,on_stderr).

    >>> colout.__args_dirty__(["colout","pattern"],"usage")
    ('pattern', 'red', 'normal', False)
    >>> colout.__args_dirty__(["colout","pattern","colors","styles"],"usage")
    ('pattern', 'colors', 'styles', False)
    >>> colout.__args_dirty__(["colout","pattern","colors","styles","True"],"usage")
    ('pattern', 'colors', 'styles', True)
    """
    import sys

    # Use a dirty argument picker
    # Check for bad usage or an help flag 
    if    len(argv) < 2 \
       or len(argv) > 9 \
       or argv[1] == "--help" \
       or argv[1] == "-h":
        print(usage+"\n")
        print("Usage:",argv[0],"<pattern> <color(s)> [<style(s)>] [<print on stderr?>] [<iterate over groups?>]")
        print("\tAvailable colors:"," ".join(colors))
        print("\tAvailable styles:"," ".join(styles))
        print("Example:",argv[0],"'^(def)\s+(\w*).*$' blue,magenta italic,bold < colout.py")
        sys.exit(1)

    assert( len(argv) >= 2 )
    # Get mandatory arguments
    pattern = argv[1]

    # default values for optional args
    color = "red"
    style = "normal"
    on_stderr = False

    if len(argv) >= 3:
        color = argv[2]
        if len(argv) >= 4:
            style = argv[3]
            if len(argv) == 5:
                on_stderr = bool(argv[4])
                if len(argv) == 6:
                    on_groups = bool(argv[5])
                    if len(argv) == 7:
                        as_colormap = bool(argv[6])
                        if len(argv) == 8:
                            as_theme = bool(argv[7])
                            if len(argv) == 9:
                                as_source = bool(argv[8])

    return pattern,color,style,on_stderr,on_groups,as_colormap,as_theme,as_source


def __args_parse__(argv,usage=""):
    """
    Parse command line arguments with the argparse library.
    Returns a tuple of (pattern,color,style,on_stderr).
    """
    parser = argparse.ArgumentParser(
        description=usage)

    parser.add_argument("pattern", metavar="REGEX", type=str, nargs=1,
            help="A regular expression")

    parser.add_argument("color", metavar="COLOR", type=str, nargs='?',
            default="red",
            help="A number in [0…255], one of the available colors or a comma-separated list of values. \
                        Available colors: "+" ".join(colors) )

    parser.add_argument("style", metavar="STYLE", type=str, nargs='?',
            default="bold",
            help="One of the available styles or a comma-separated list of styles.\
                        Available styles: "+" ".join(styles) )

    parser.add_argument("-e", "--stderr", action="store_true",
            help="Output on the stderr instead of stdout")

    parser.add_argument("-g", "--groups", action="store_true",
            help="For color maps (random, rainbow), iterate over matching groups in the pattern instead of over patterns")

    parser.add_argument("-c", "--colormap", action="store_true",
            help="Use the given colors as a colormap (cycle the colors at each match)")

    parser.add_argument("-t", "--theme", action="store_true",
            help="Interpret REGEX as a theme")

    parser.add_argument("-s", "--source", action="store_true",
            help="""Interpret REGEX as a source code readable by the Pygments library.
            If the first letter of PATTERN is upper case, use the 256 colors mode,
            if it is lower case, use the 8 colors mode.
            In 256 colors, interpret COLOR as a Pygments style (e.g. "default").""")

    args = parser.parse_args()

    return args.pattern[0], args.color, args.style, args.stderr, args.groups, args.colormap, args.theme, args.source


def write( colored, on_stderr=False ):
    """
    If on_stderr, write "colored" on sys.stderr, else write it on sys.stdout.
    Then flush.
    """
    if on_stderr:
        sys.stderr.write(colored)
        sys.stderr.flush()
    else:
        sys.stdout.write(colored)
        sys.stdout.flush()


if __name__ == "__main__":
    import sys

    usage="A regular expression based formatter that color up an arbitrary text stream."

    try:
        import argparse

    # if argparse is not installed
    except ImportError:
        pattern,color,style,on_stderr,on_groups,as_colormap,as_theme,as_source = __args_dirty__(sys.argv,usage)

    # if argparse is available
    else:
        pattern,color,style,on_stderr,on_groups,as_colormap,as_theme,as_source = __args_parse__(sys.argv,usage)

    # use the generator: output lines as they come
    if as_colormap == True and color != "rainbow":
        colormap = color.split(",")  # replace the colormap by the given colors
        color = "colormap"  # use the keyword to switch to colormap instead of list of colors

    # load available themes
    themes = {}
    import glob
    for f in glob.iglob("colout_*.py"):
        module = ".".join(f.split(".")[:-1])
        name = "_".join(module.split("_")[1:])
        themes[name] = __import__(module)

    # load available pygments lexers
    try:
        from pygments.lexers import get_all_lexers
        from pygments.lexers import get_lexer_by_name
        from pygments import highlight
        from pygments.formatters import Terminal256Formatter
        from pygments.formatters import TerminalFormatter
    except ImportError:
        pass
    else:
        lexers = []
        for lexer in get_all_lexers():
            lexers.append( lexer[1][0] )

    # if theme
    if as_theme:
        assert( pattern in themes.keys() )
        while True:
            try:
                item = sys.stdin.readline()
            except KeyboardInterrupt:
                break
            if not item:
                break
            colored = themes[pattern].theme(item)
            write(colored)

    # if pygments
    elif as_source:
        assert( pattern.lower() in lexers )
        lexer = get_lexer_by_name(pattern.lower())
        # Python => 256 colors, python => 8 colors
        ask_256 = pattern[0].isupper()
        if ask_256:
            try:
                formatter = Terminal256Formatter(style = color)
            except: # style not found
                formatter = Terminal256Formatter()
        else:
            formatter = TerminalFormatter()

        while True:
            try:
                item = sys.stdin.readline()
            except KeyboardInterrupt:
                break
            if not item:
                break
            colored = highlight( item, lexer, formatter)
            write(colored)

    # if color
    else:
        for colored in colorgen( sys.stdin, pattern, color, style, on_groups):
            write(colored)