color.py

# Copyright (C) 2008 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys

import pager

COLORS = {
    None: -1,
    "normal": -1,
    "black": 0,
    "red": 1,
    "green": 2,
    "yellow": 3,
    "blue": 4,
    "magenta": 5,
    "cyan": 6,
    "white": 7,
}

ATTRS = {None: -1, "bold": 1, "dim": 2, "ul": 4, "blink": 5, "reverse": 7}

RESET = "\033[m"


def is_color(s):
    return s in COLORS


def is_attr(s):
    return s in ATTRS


def _Color(fg=None, bg=None, attr=None):
    fg = COLORS[fg]
    bg = COLORS[bg]
    attr = ATTRS[attr]

    if attr >= 0 or fg >= 0 or bg >= 0:
        need_sep = False
        code = "\033["

        if attr >= 0:
            code += chr(ord("0") + attr)
            need_sep = True

        if fg >= 0:
            if need_sep:
                code += ";"
            need_sep = True

            if fg < 8:
                code += "3%c" % (ord("0") + fg)
            else:
                code += "38;5;%d" % fg

        if bg >= 0:
            if need_sep:
                code += ";"

            if bg < 8:
                code += "4%c" % (ord("0") + bg)
            else:
                code += "48;5;%d" % bg
        code += "m"
    else:
        code = ""
    return code


DEFAULT = None


def SetDefaultColoring(state):
    """Set coloring behavior to |state|.

    This is useful for overriding config options via the command line.
    """
    if state is None:
        # Leave it alone -- return quick!
        return

    global DEFAULT
    state = state.lower()
    if state in ("auto",):
        DEFAULT = state
    elif state in ("always", "yes", "true", True):
        DEFAULT = "always"
    elif state in ("never", "no", "false", False):
        DEFAULT = "never"


class Coloring(object):
    def __init__(self, config, section_type):
        self._section = "color.%s" % section_type
        self._config = config
        self._out = sys.stdout

        on = DEFAULT
        if on is None:
            on = self._config.GetString(self._section)
            if on is None:
                on = self._config.GetString("color.ui")

        if on == "auto":
            if pager.active or os.isatty(1):
                self._on = True
            else:
                self._on = False
        elif on in ("true", "always"):
            self._on = True
        else:
            self._on = False

    def redirect(self, out):
        self._out = out

    @property
    def is_on(self):
        return self._on

    def write(self, fmt, *args):
        self._out.write(fmt % args)

    def flush(self):
        self._out.flush()

    def nl(self):
        self._out.write("\n")

    def printer(self, opt=None, fg=None, bg=None, attr=None):
        s = self
        c = self.colorer(opt, fg, bg, attr)

        def f(fmt, *args):
            s._out.write(c(fmt, *args))

        return f

    def nofmt_printer(self, opt=None, fg=None, bg=None, attr=None):
        s = self
        c = self.nofmt_colorer(opt, fg, bg, attr)

        def f(fmt):
            s._out.write(c(fmt))

        return f

    def colorer(self, opt=None, fg=None, bg=None, attr=None):
        if self._on:
            c = self._parse(opt, fg, bg, attr)

            def f(fmt, *args):
                output = fmt % args
                return "".join([c, output, RESET])

            return f
        else:

            def f(fmt, *args):
                return fmt % args

            return f

    def nofmt_colorer(self, opt=None, fg=None, bg=None, attr=None):
        if self._on:
            c = self._parse(opt, fg, bg, attr)

            def f(fmt):
                return "".join([c, fmt, RESET])

            return f
        else:

            def f(fmt):
                return fmt

            return f

    def _parse(self, opt, fg, bg, attr):
        if not opt:
            return _Color(fg, bg, attr)

        v = self._config.GetString("%s.%s" % (self._section, opt))
        if v is None:
            return _Color(fg, bg, attr)

        v = v.strip().lower()
        if v == "reset":
            return RESET
        elif v == "":
            return _Color(fg, bg, attr)

        have_fg = False
        for a in v.split(" "):
            if is_color(a):
                if have_fg:
                    bg = a
                else:
                    fg = a
            elif is_attr(a):
                attr = a

        return _Color(fg, bg, attr)