colors.py

# colors.py

import copy
import platform
import sys

# basic color definitions
BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)

# default values
_defaults = {
    'foreground' : None,
    'background' : None,
    'bold' : False,
}

class RangeSet(object):
    """An object representing the set [start, end), or possibly a
    finite union of such sets."""
    def __init__(self, *args):
        super(RangeSet, self).__init__()
        if len(args) == 2:
            start, end = args
            if start > end:
                raise ValueError("cannot create a range set of negative size")
            elif start == end:
                self.ranges = []
            else:
                self.ranges = [(start, end)]
        elif len(args) == 1:
            self.ranges = args[0]
            self.ranges.sort()
        else:
            raise TypeError("__init__() takes 2 or 3 arguments (%i given)" % (len(args),))

    def __repr__(self):
        if not self.ranges:
            return "{null set}"
        return " U ".join(["[%i, %i)" % r for r in self.ranges])
    
    @property
    def min(self):
        if not self.ranges:
            return None
        return self.ranges[0][0]
    @property
    def max(self):
        if not self.ranges:
            return None
        return self.ranges[-1][1]
    @property
    def empty(self):
        if not self.ranges:
            return True
        return False
    
    def contains(self, i):
        for r in self.ranges:
            if i >= r[0] and i < r[1]:
                return True
        return False
    
    def translate(self, i):
        ranges = [(i + r[0], i + r[1]) for r in self.ranges]
        return type(self)(ranges)
    
    # this "merge" and the functions based off it are ridiculously
    # inefficient on large sets -- they could use optimizations
    
    def merge(self, test, *args):
        if len(args) > 1:
            other = type(self)(*args)
        else:
            other = args[0]
        
        ranges = []
        start_index = 0
        in_range = False
        for i in xrange(min(self.min, other.min), max(self.max, other.max) + 1):
            if not in_range and test(self.contains(i), other.contains(i)):
                in_range = True
                start_index = i
            elif in_range and not test(self.contains(i), other.contains(i)):
                in_range = False
                ranges.append((start_index, i))
        return type(self)(ranges)
    
    def union(self, *args):
        return self.merge(lambda a, b: a or b, *args)
    def intersection(self, *args):
        return self.merge(lambda a, b: a and b, *args)
    def subtract(self, *args):
        return self.merge(lambda a, b: a and not b, *args)
    
class ColoredImmutableSequence(object):
    __color_type__ = None
    
    def __init__(self, *args, **kwargs):
        super(ColoredImmutableSequence, self).__init__()
        self.__colors__ = {}
        for prop in _defaults:
            self.__colors__[prop] = {_defaults[prop] : RangeSet(0, len(self))}
        self.setcolor(**kwargs)
    
    def __add__(self, rhs):
        result = self.__color_type__.__add__(self, rhs)
        result = type(self)(result)
        result.__colors__ = copy.deepcopy(self.__colors__)
        
        if isinstance(rhs, ColoredImmutableSequence):
            for prop in rhs.__colors__:
                for val in rhs.__colors__[prop]:
                    r = rhs.__colors__[prop][val]
                    kwargs = {prop : val}
                    result.setcolor(r.translate(len(self)), **kwargs)
        
        return result
    
    def __radd__(self, lhs):
        result = lhs.__add__(self)
        result = type(self)(result)
        
        for prop in self.__colors__:
            for val in self.__colors__[prop]:
                r = self.__colors__[prop][val]
                kwargs = {prop : val}
                result.setcolor(r.translate(len(lhs)), **kwargs)
        
        return result
    
    def setcolor(self, *args, **kwargs):
        if len(args) == 0:
            r = RangeSet(0, len(self))
        elif len(args) == 1:
            r = args[0]
        else:
            r = RangeSet(*args)
        
        for prop in kwargs:
            val = kwargs[prop]
            if not prop in self.__colors__:
                self.__colors__[prop] = {}
            data = self.__colors__[prop]
            if not val in data:
                data[val] = r
            data_to_remove = []
            for iterval in data:
                if iterval == val:
                    data[iterval] = data[iterval].union(r)
                else:
                    data[iterval] = data[iterval].subtract(r)
                    if data[iterval].empty:
                        data_to_remove.append(iterval)
            for iterval in data_to_remove:
                del data[iterval]
        
        return self
    
    # this splitter function is inefficient on large strings, could be improved
    def splitcolors(self):
        last_index = 0
        last_settings = {}
        ranges = []
        for i in xrange(len(self) + 1):
            settings = copy.deepcopy(_defaults)
            for prop in self.__colors__:
                for val in self.__colors__[prop]:
                    valrange = self.__colors__[prop][val]
                    if valrange.contains(i):
                        settings[prop] = val
            if last_index == i:
                last_settings = settings
            elif settings != last_settings or i == len(self):
                ranges.append((last_index, i, last_settings))
                last_index = i
                last_settings = settings
        return [(self[r[0]:r[1]], r[2]) for r in ranges]

class ColoredString(ColoredImmutableSequence, str):
    __color_type__ = str
    
    def __str__(self):
        return self
    def __unicode__(self):
        return self

# convenience functions
def colored(s, **kwargs):
    if not isinstance(s, ColoredImmutableSequence):
        s = ColoredString(s)
    return s.setcolor(**kwargs)

def black(s): return colored(s, foreground=BLACK)
def red(s): return colored(s, foreground=RED)
def green(s): return colored(s, foreground=GREEN)
def yellow(s): return colored(s, foreground=YELLOW)
def blue(s): return colored(s, foreground=BLUE)
def magenta(s): return colored(s, foreground=MAGENTA)
def cyan(s): return colored(s, foreground=CYAN)
def white(s): return colored(s, foreground=WHITE)

def bgblack(s): return colored(s, background=BLACK)
def bgred(s): return colored(s, background=RED)
def bggreen(s): return colored(s, background=GREEN)
def bgyellow(s): return colored(s, background=YELLOW)
def bgblue(s): return colored(s, background=BLUE)
def bgmagenta(s): return colored(s, background=MAGENTA)
def bgcyan(s): return colored(s, background=CYAN)
def bgwhite(s): return colored(s, background=WHITE)

def bold(s): return colored(s, bold=True)

def reset(s): return colored(s, **_defaults)

# output handlers

class ColoredOutput(object):
    def __init__(self, target):
        super(ColoredOutput, self).__init__()
        self.target = target
        
    def write(self, s):
        if isinstance(s, ColoredImmutableSequence):
            old_colors = _defaults
            for s, colors in s.splitcolors():
                colordiff = {}
                for prop in colors:
                    if not prop in old_colors or old_colors[prop] != colors[prop]:
                        colordiff[prop] = colors[prop]
                for prop in old_colors:
                    if not prop in colors:
                        colordiff[prop] = None
                self.colors_set(colordiff)
                self.target.write(s)
                
                old_colors = colors
            self.colors_reset()
        else:
            self.target.write(s)
    
    # dummy, do-nothing implementations
    def colors_set(self, colors):
        raise NotImplemented("colors_set")
    def colors_reset(self):
        raise NotImplemented("colors_reset")
    
    # pass-through all other getattr, setattr calls
    def __getattr__(self, attr):
        return getattr(self.target, attr)
    def __setattr__(self, attr, val):
        if attr == "target":
            self.__dict__[attr] = val
        else:
            setattr(self.target, attr, val)

class ANSIOutput(ColoredOutput):
    def colors_set(self, colors):
        if 'foreground' in colors:
            if colors['foreground']:
                self.target.write("\033[%im" % (30 + colors['foreground'],))
            else:
                self.target.write("\033[39m")
        if 'background' in colors:
            if colors['background']:
                self.target.write("\033[%im" % (40 + colors['background'],))
            else:
                self.target.write("\033[49m")
        if 'bold' in colors:
            if colors['bold']:
                self.target.write("\033[1m")
            else:
                self.target.write("\033[22m")
    def colors_reset(self):
        self.target.write("\033[0m")

# some logic to wrap a file automatically, using color if appropriate
def wrap(f):
    if platform.system() == 'Windows':
        # do nothing, for now
        return f
    elif hasattr(f, 'isatty') and f.isatty():
        return ANSIOutput(f)
    
    # nothing matched, default to no color
    return f

# auto-wrap stdout and stderr
sys.stdout = wrap(sys.stdout)
sys.stderr = wrap(sys.stderr)

if __name__ == "__main__":
    import random
    
    def complex(s, foreground=MAGENTA):
        s = colored(s, foreground=foreground)
        for i in range(3):
            s.setcolor(2*i, 2*i + 1, bold=True)
        return s
    
    def complicated(s):
        s = colored(s)
        colors = range(8)
        bolds = [False, True]
        for i in range(len(s)):
            s.setcolor(i, i + 1, foreground=random.choice(colors), background=random.choice(colors), bold=random.choice(bolds))
        return s
    
    def rainbow(s):
        s = ColoredString(s)
        colors = [MAGENTA, RED, YELLOW, GREEN, CYAN, BLUE]
        for i in range(len(s)):
            s.setcolor(i, i + 1, foreground=colors[i % len(colors)])
        return bold(s)
    
    print
    print green("The " + bold("Zen") + " of Python") + ", by Tim Peters"
    print
    print cyan("Beautiful") + " is better than " + bgyellow("ugly") + "."
    print bold(red("Explicit")) + " is better than " + white("implicit") + "."
    print magenta("Simple") + " is better than " + complex("complex") + "."
    print complex("Complex") + " is better than " + complicated("complicated") + "."
    print white("Flat") + " is better than " + rainbow("nested") + "."
    print blue(bold("Sparse")) + " is better than " + bgblue("dense") + "."
    print "Readability counts."
    print green("Special cases") + " aren't " + green("special") + " enough to " + red("break the rules") + "."
    print "Although practicality beats " + cyan("purity") + "."
    print red("Errors") + " should never pass " + white("silently") + "."
    print white("Unless explicitly silenced.")
    print "In the face of " + green("ambiguity") + ", refuse the " + yellow("temptation to guess") + "."
    print "There should be " + cyan("one") + "-- and preferably " + magenta(bold("only")) + cyan(" one") + " --" + yellow(bold("obvious")) + " way to do it."
    print "Although that way may not be " + bold(yellow("obvious")) + " at first unless you're " + bold(red("Du") + "t" + blue("ch")) + "."
    print green("Now") + " is better than " + bold(red("never")) + "."
    print "Although " + bold(red("never")) + " is often better than " + cyan("*right*") + " " + green("now") + "."
    print "If the implementation is " + red(bold("hard to explain")) + ", it's a " + red("bad idea") + "."
    print "If the implementation is " + cyan(bold("easy to explain")) + ", it may be a " + cyan("good idea") + "."
    print green(bold("Namespaces")) + " are one honking " + magenta(bold("great idea")) + " -- let's do more of those!"
    print