Color remapping

#   Copyright 2013 Adam Ehlers Nyholm Thomsen
#   Copyright 2012 Darsh Ranjan
#
#   This file is based on code from Darsh Ranjan's termcolors project.
#
#   termcolors is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   termcolors is distributed in the hope that it will be useful, but
#   WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#   General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with termcolors.  If not, see
#   <http://www.gnu.org/licenses/>.

import os
import re
import select
import termios
from sys import version_info

#######################################################################
# Query-related error conditions

class TerminalQueryError(Exception):

    '''
    Base class for the other exceptions.

    '''

    def __init__(self, message):
        Exception.__init__(self, message)


class TerminalSetupError(TerminalQueryError):

    '''
    We couldn't set up the terminal properly.

    '''

    def __init__(self, fd):
        TerminalQueryError.__init__(
            self,
            ("Couldn't set up terminal on file " +
             ("descriptor %d" % fd)))


class NoResponseError(TerminalQueryError):

    '''
    The terminal didn't respond, or we were too impatient.

    '''

    def __init__(self, q):
        TerminalQueryError.__init__(
            self,
            "Timeout on query " + repr(q))


class TerminalUninitializedError(TerminalQueryError):

    '''
    Someone tried to do something without setting up the terminal
    properly (by calling TerminalQueryContext.__enter__).

    '''

    def __init__(self, fd):
        TerminalQueryError.__init__(
            self,
            (("Terminal on file descriptor %d " % fd) +
             "not set up"))


########################################################################

class TerminalQueryContext(object):

    '''
    Context manager for terminal RGB queries.

    '''

    # Operating system command
    osc = "\033]"

    # String terminator
    #  ("\033\\" is another option, but "\007" seems to be understood by
    #  more terminals.  Terminology, for example, doesn't seem to like
    #  "\033\\".)
    st = "\007"

    # Control sequence introducer
    csi = "\033["

    # ANSI SGR0
    reset = csi + 'm'


    def __init__(self, fd=0, timeout = -1):
        '''
        fd: open file descriptor referring to the terminal we care
        about.  The default (0) is almost always correct.

        timeout: how many milliseconds to wait for a response, a negative
            number meaning "infinite". A larger timeout is safer but will cause
            this function to take proportionally more time.
        '''
        self.tc_save = None
        self.fd = fd

        self.num_errors = 0
        self.timeout = timeout


    def __enter__(self):
        '''
        Set up the terminal for queries.

        '''
        self.tc_save = termios.tcgetattr(self.fd)

        tc = termios.tcgetattr(self.fd)

        # Don't echo the terminal's responses
        tc[3] &= ~termios.ECHO

        # Noncanonical mode (i.e., disable buffering on the terminal
        # level)
        tc[3] &= ~termios.ICANON

        # Make input non-blocking
        tc[6][termios.VMIN] = 0
        tc[6][termios.VTIME] = 0

        termios.tcsetattr(self.fd, termios.TCSANOW, tc)

        # Check if it succeeded
        if termios.tcgetattr(self.fd) != tc:
            termios.tcsetattr(self.fd, termios.TCSANOW, self.tc_save)
            raise TerminalSetupError(self.fd)

        self.P = select.poll()
        self.P.register(self.fd, select.POLLIN)

        return self


    def __exit__(self, exc_type, exc_value, traceback):
        '''
        Reset the terminal to its original state.

        '''
        self.flush_input()

        if self.tc_save is not None:
            termios.tcsetattr(self.fd, termios.TCSANOW, self.tc_save)

        del self.P


    def get_num_colors(self):
        '''
        Attempt to determine the number of colors we are able to query
        from the terminal.

        '''
        # We won't count failed queries in this function, since we're
        # guaranteed to fail a few.
        num_errors = self.num_errors

        if not self.get_indexed_color(0):
            return 0

        a = 0
        b = 1
        while self.get_indexed_color(b):
            a = b
            b += b

        while b - a > 1:
            c = (a + b)>>1
            if self.get_indexed_color(c):
                a = c
            else:
                b = c

        self.num_errors = num_errors
        return b


    def get_all_indexed_colors(self, limit=-1):
        '''
        Query as many indexed colors as possible up to `limit'
        and return them all in a list.
        '''
        num_colors = self.get_num_colors()
        if limit >= 0:
            num_colors = min(num_colors, limit)
        if num_colors == 0:
            return []

        return self.color_query(range(num_colors))

    # Wrappers for xterm & urxvt operating system controls.
    #
    # These codes are all common to xterm and urxvt. Their responses
    # aren't always in the same format (xterm generally being more
    # consistent), but the regular expression used to parse the
    # responses is general enough to work for both.
    #
    # Note: none of these functions is remotely thread-safe.


    def get_indexed_color(self, a):
        '''
        Get color number `a'.

        '''
        return self.color_query([a,])


    def flush_input(self):
        '''
        Discard any input that can be read at this moment.

        '''
        while self.P.poll(0):
            os.read(self.fd, 4096)

    # Patterns matching unsigned decimal and hexadecimal integer
    # literals
    ndec = "[0-9]+"
    nhex = "[0-9a-fA-F]+"

    # The "guard" query and its response pattern.
    q_guard = csi + "6n"

    str_guard = "(.*)\033\\[{ndec};{ndec}R".format(**vars())
    re_guard = re.compile(str_guard)

    # This is what we expect the terminal's response to a query for a
    # color to look like.  If we didn't care about urxvt, we could get
    # away with a simpler implementation here, since xterm and vte seem
    # to give pretty consistent and systematic responses.  But I
    # actually use urxvt most of the time, so....
    str_color = "\033\\]4;(?:(?P<idx>{ndec});)?(?P<color>[^;\007]+)\007".format(**vars())
    re_color = re.compile(str_color)

    # This is what the resp
    str_default_color = "\033\\](?P<key_id>{ndec});(?P<color>[^;\007]+)\007".format(**vars())
    re_default_color = re.compile(str_default_color)

    def color_query(self, q):
        '''
        Query color-valued terminal parameter.

        Arguments:
            q: The query code as a sequence of nonnegative integers,
                i.e., [q0, q1, ...] if the escape sequence in
                pseudo-Python is

                    "\033]4;{q0};?;{q1};?;...\007"

        Return:
            A list of tuples (idx, color) where color is a string representing
            the color value of the color with index idx.  I.e. the response
            should look like [(q0, color0), (q1, color1), ...].  If the
            terminal provides an unparseable (or no) response, then None will
            be returned.

        See
            http://invisible-island.net/xterm/ctlseqs/ctlseqs.html

        ("Operating System Controls") to see the various queries
        supported by xterm.  Urxvt supports some, but not all, of them,
        and has a number of its own (see man -s7 urxvt).

        self.__enter__ must be called prior to calling this function, or
        TerminalUninitializedError will be raised.

        '''
        if not q:
            return None

        query = (self.osc + '4;' +
                 ';'.join(('%d;?' % k for k in q)) +
                 self.st)

        try:
            response = self.guarded_query(query)
        except NoResponseError:
            return None

        res = []
        for i, m in zip(q, self.re_color.finditer(response)):
            idx = m.group('idx')
            if idx != None:
                assert int(idx) == i  # Make sure we get the responses in the correct order
            res.append((i, m.group('color')))

        # Since we do multiple queries at once, try to calculate the number of
        # errors.
        num_errors = len(q) - len(res)
        assert num_errors >= 0
        self.num_errors += num_errors

        return res

    def color_setup(self, q):
        """Setup color values of the terminal.

        Arguments:
            q:  A list of tuples (idx, color) such that color is a string
                representing the color value that idx should take.  Which
                strings are permittable depends on the terminal, but most
                terminals accept anything you can throw at XParseColor.
                Specifically "#rrggbb" works.

                If q = [(i0, col0), (i1, col1), ...] then the escape sequence
                used is:
                    "\033]4;{i0};{col0};{i1};{col1};...\007"

        Returns:  None

        See
            http://invisible-island.net/xterm/ctlseqs/ctlseqs.html

        ("Operating System Controls") to see the various queries
        supported by xterm.  Urxvt supports some, but not all, of them,
        and has a number of its own (see man -s7 urxvt).

        self.__enter__ must be called prior to calling this function, or
        TerminalUninitializedError will be raised.
        """
        if not q:
            return None

        query = (self.osc + '4;' +
                 ';'.join(('%d;%s' % k for k in q)) +
                 self.st)

        try:
            response = self.guarded_query(query)
        except NoResponseError:
            return None

    def get_default_color(self, key_id):
        """Gets the values of the default color with id key_id

        The following key_id's of interest(it has more) are specified in the
        xterm ctlseqs.txt file:
            10     Foreground color
            11     Background color
            12     Text cursor color
            13     Mouse foreground color
            14     Mouse background color

        Note this is NOT supported on VTE, and I have only found it to work on
        xterm and urxvt"""

        query = (self.osc + ('%d;?' % key_id) + self.st)

        try:
            response = self.guarded_query(query)
        except NoResponseError:
            return None

        m = self.re_default_color.match(response)
        assert int(m.group('key_id')) == key_id
        return m.group('color')

    def set_default_color(self, key_id, color):
        """Sets the values of the default color with key_id

        Note this is NOT supported on VTE, and I have only found it to work on
        xterm and urxvt"""
        query = (self.osc + ('%d;%s' % (key_id,color)) + self.st)
        try:
            response = self.guarded_query(query)
        except NoResponseError:
            return None

    # If a terminal sees an escape sequence it doesn't like, it will
    # simply ignore it.  Also, it's hard to predict how long a terminal
    # will take to respond to a query it does like.  However, some
    # escape sequences, like "\033[6n", will produce a predictable
    # response on *most* (but not all) terminals, and this fact can be
    # used to test for the absence of a response to a particular query
    # on such terminals.

    def guarded_query(self, q, flush=True):
        '''
        Send the terminal query string `q' and return the terminal's
        response.

        Arguments:
            q: the query string to send to the terminal.

            flush: whether to discard waiting input before sending the
                query.  It usually makes sense to do this, but note that
                the terminal may still send seemingly unsolicited data
                (possibly in response to an earlier query) after the
                input is flushed, so flushing the input does not
                provide any guarantees.

        Return: The terminal's response to the query as a string.

        Errors:
            NoResponseError will be raised if the query times out.

            TerminalUninitializedError will be raised if this instance's
            context has not been __enter__-ed.

        '''
        if not hasattr(self, "P"):
            raise TerminalUninitializedError(self.fd)

        query = q + self.q_guard

        if flush:
            self.flush_input()

        os.write(self.fd, query.encode())

        response = ""

        while self.P.poll(self.timeout):
            s = os.read(self.fd, 4096)
            if version_info.major >= 3:
                s = s.decode()
            response += s

            m = self.re_guard.match(response)

            if m:
                return m.group(1)
        else:
            self.num_errors += 1
            raise NoResponseError(query)

def get_random_rgb():
    import random
    return '#' + ''.join(('%x' % random.randint(0,15) for i in range(6)))

def get_light_name(name):
    """Get the name of the corresponding color in the light palette"""
    if name.startswith('base'):
        # The base colors are opposite in the light palette
        if len(name) == 6:
            return 'base' + name[5]
        elif len(name) == 5:
            return 'base0' + name[4]
        else:
            assert False
    else:
        return name

def get_light(name):
    """Get the information for the corresponding light color"""
    light_val = sol[get_light_name(name)] # The light color information, for the hex and the 256 index
    dark_val = sol[name] # The dark color information, for the 16 and 8 indices
    # Return the light hex, dark 16 index, dark 8 indices, and light 256 color
    # We return the light 256 color, since the 256 color map is supposed to
    # approximate rgb colors and hence switching between light and dark
    # mode should be done in the application color scheme not in the
    # terminal color map.
    return (light_val[0], dark_val[1], dark_val[2], light_val[3])

def get_colors(mode):
    """Get the index to hex mapping for the given color mode.

    Arguments:
        mode:
            Get the index to hex mapping for mode, currently supported modes are:
                16-dark     The dark version of the 16-color palette, for legacy applications like irssi
                16-light    The light version of the 16-color palette, for legacy applications like irssi
                256         The approximate 256 color entries for xterm and the exact solarized hex values

    Returns:
        A list of tuples (idx, color) where idx is an integer and color is a
        hex RGB representation of the color.
    """

    return modes[mode]

def get_default_colors_aux(light = False):
    get_hex_dark = lambda name: sol[name][0]
    get_hex_light = lambda name: get_light(name)[0]
    get_color = get_hex_light if light else get_hex_dark

    return [(i, get_color(name)) for (i, name) in default_colors.items()]

def get_default_colors(mode):
    if not mode:
        return []
    return get_default_colors_aux(mode == 'light')

def get_default_ids():
    return list(default_colors.keys())

#SOLARIZED HEX     16/8 TERMCOL  XTERM/HEX   L*A*B      RGB         HSB
#--------- ------- ---- -------  ----------- ---------- ----------- -----------
#base03    #002b36  8/4 brblack  234 #1c1c1c 15 -12 -12   0  43  54 193 100  21
#base02    #073642  0/4 black    235 #262626 20 -12 -12   7  54  66 192  90  26
#base01    #586e75 10/7 brgreen  240 #585858 45 -07 -07  88 110 117 194  25  46
#base00    #657b83 11/7 bryellow 241 #626262 50 -07 -07 101 123 131 195  23  51
#base0     #839496 12/6 brblue   244 #808080 60 -06 -03 131 148 150 186  13  59
#base1     #93a1a1 14/4 brcyan   245 #8a8a8a 65 -05 -02 147 161 161 180   9  63
#base2     #eee8d5  7/7 white    254 #e4e4e4 92 -00  10 238 232 213  44  11  93
#base3     #fdf6e3 15/7 brwhite  230 #ffffd7 97  00  10 253 246 227  44  10  99
#yellow    #b58900  3/3 yellow   136 #af8700 60  10  65 181 137   0  45 100  71
#orange    #cb4b16  9/3 brred    166 #d75f00 50  50  55 203  75  22  18  89  80
#red       #dc322f  1/1 red      160 #d70000 50  65  45 220  50  47   1  79  86
#magenta   #d33682  5/5 magenta  125 #af005f 50  65 -05 211  54 130 331  74  83
#violet    #6c71c4 13/5 brmagenta 61 #5f5faf 50  15 -45 108 113 196 237  45  77
#blue      #268bd2  4/4 blue      33 #0087ff 55 -10 -45  38 139 210 205  82  82
#cyan      #2aa198  6/6 cyan      37 #00afaf 60 -35 -05  42 161 152 175  74  63
#green     #859900  2/2 green     64 #5f8700 60 -20  65 133 153   0  68 100  60

sol = {
    # Name    : (HEX,      16, 8, 256)
    'base03'  : ('#002b36',  8, 4, 234),
    'base02'  : ('#073642',  0, 4, 235),
    'base01'  : ('#586e75', 10, 7, 240),
    'base00'  : ('#657b83', 11, 7, 241),
    'base0'   : ('#839496', 12, 6, 244),
    'base1'   : ('#93a1a1', 14, 4, 245),
    'base2'   : ('#eee8d5',  7, 7, 254),
    'base3'   : ('#fdf6e3', 15, 7, 230),
    'yellow'  : ('#b58900',  3, 3, 136),
    'orange'  : ('#cb4b16',  9, 3, 166),
    'red'     : ('#dc322f',  1, 1, 160),
    'magenta' : ('#d33682',  5, 5, 125),
    'violet'  : ('#6c71c4', 13, 5,  61),
    'blue'    : ('#268bd2',  4, 4,  33),
    'cyan'    : ('#2aa198',  6, 6,  37),
    'green'   : ('#859900',  2, 2,  64),
}

default_colors = {
    10 : 'base0',  #Foreground
    11 : 'base03', #Background
    12 : 'base1',  #Text cursor color
#    13 : 'base1',  #Mouse foreground color
#    14 : 'base01', #Mouse background color
}

modes = {
    '16-dark'  : [(i16,  h) for (h, i16, i8, i256) in sol.values()],
    '16-light' : [(i16,  h) for (h, i16, i8, i256) in map(get_light, sol)],
    '256'      : [(i256, h) for (h, i16, i8, i256) in sol.values()],
}

def construct_argparser():
    import argparse
    parser = argparse.ArgumentParser(description = 'Color remapping for xterm-like terminals.')
    parser.add_argument('-n', '--no-reset', dest = 'reset', action = 'store_const', const = False, default = True, help = 'Don\'t reset the colors afterwards(keep solarized colors)')
    parser.add_argument('-m', '--mode', dest = 'mode', action = 'store', default = '256', help = 'The color palette to map', choices = ['16-dark', '16-light', '256'])
    parser.add_argument('-d', '--set-defaults', dest = 'default_cols', action = 'store', default = False, help = 'Set the default colors(foreground, background, cursor, etc.)', choices = ['dark', 'light'])
    parser.add_argument('CMD', nargs=argparse.REMAINDER, help = 'An (optional) command line application with arguments to run in the remapped color environment, for when you don\'t pass --no-reset')
    return parser

if __name__ == '__main__':
    import subprocess, sys
    parser = construct_argparser()
    args = parser.parse_args()
    with TerminalQueryContext() as tq:
        def set_defaults(cols):
            for i, col in cols:
                tq.set_default_color(i, col)

        orig = [] # The original color mappings
        orig_def = []

        new_def = get_default_colors(args.default_cols)
        m = get_colors(args.mode)
        if args.reset:
            orig = tq.color_query([idx for (idx, col) in m])
            orig_def = [(i, tq.get_default_color(i)) for i, col in new_def]

        tq.color_setup(m)
        set_defaults(new_def)

        ret = 0
        try:
            if args.CMD:
                ret = subprocess.call(args.CMD)
        finally:
            tq.color_setup(orig)
            set_defaults(orig_def)
        if False: #Debug output
            print('Errors: ', tq.num_errors)
        sys.exit(ret)