path: root/scripts/lib/wic/msger.py

#!/usr/bin/env python -tt
# vim: ai ts=4 sts=4 et sw=4
#
# Copyright (c) 2009, 2010, 2011 Intel, Inc.
#
# This program 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; version 2 of the License
#
# This program 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 this program; if not, write to the Free Software Foundation, Inc., 59
# Temple Place - Suite 330, Boston, MA 02111-1307, USA.

import os
import sys
import re
import time

__ALL__ = ['set_mode',
           'get_loglevel',
           'set_loglevel',
           'set_logfile',
           'raw',
           'debug',
           'verbose',
           'info',
           'warning',
           'error',
           'ask',
           'pause',
          ]

# COLORs in ANSI
INFO_COLOR = 32 # green
WARN_COLOR = 33 # yellow
ERR_COLOR = 31 # red
ASK_COLOR = 34 # blue
NO_COLOR = 0

PREFIX_RE = re.compile('^<(.*?)>\s*(.*)', re.S)

INTERACTIVE = True

LOG_LEVEL = 1
LOG_LEVELS = {
    'quiet': 0,
    'normal': 1,
    'verbose': 2,
    'debug': 3,
    'never': 4,
}

LOG_FILE_FP = None
LOG_CONTENT = ''
CATCHERR_BUFFILE_FD = -1
CATCHERR_BUFFILE_PATH = None
CATCHERR_SAVED_2 = -1

def _general_print(head, color, msg=None, stream=None, level='normal'):
    global LOG_CONTENT
    if not stream:
        stream = sys.stdout

    if LOG_LEVELS[level] > LOG_LEVEL:
        # skip
        return

    # encode raw 'unicode' str to utf8 encoded str
    if msg and isinstance(msg, unicode):
        msg = msg.encode('utf-8', 'ignore')

    errormsg = ''
    if CATCHERR_BUFFILE_FD > 0:
        size = os.lseek(CATCHERR_BUFFILE_FD, 0, os.SEEK_END)
        os.lseek(CATCHERR_BUFFILE_FD, 0, os.SEEK_SET)
        errormsg = os.read(CATCHERR_BUFFILE_FD, size)
        os.ftruncate(CATCHERR_BUFFILE_FD, 0)

    # append error msg to LOG
    if errormsg:
        LOG_CONTENT += errormsg

    # append normal msg to LOG
    save_msg = msg.strip() if msg else None
    if save_msg:
        timestr = time.strftime("[%m/%d %H:%M:%S %Z] ", time.localtime())
        LOG_CONTENT += timestr + save_msg + '\n'

    if errormsg:
        _color_print('', NO_COLOR, errormsg, stream, level)

    _color_print(head, color, msg, stream, level)

def _color_print(head, color, msg, stream, level):
    colored = True
    if color == NO_COLOR or \
       not stream.isatty() or \
       os.getenv('ANSI_COLORS_DISABLED') is not None:
        colored = False

    if head.startswith('\r'):
        # need not \n at last
        newline = False
    else:
        newline = True

    if colored:
        head = '\033[%dm%s:\033[0m ' %(color, head)
        if not newline:
            # ESC cmd to clear line
            head = '\033[2K' + head
    else:
        if head:
            head += ': '
            if head.startswith('\r'):
                head = head.lstrip()
                newline = True

    if msg is not None:
        if isinstance(msg, unicode):
            msg = msg.encode('utf8', 'ignore')

        stream.write('%s%s' % (head, msg))
        if newline:
            stream.write('\n')

    stream.flush()

def _color_perror(head, color, msg, level='normal'):
    if CATCHERR_BUFFILE_FD > 0:
        _general_print(head, color, msg, sys.stdout, level)
    else:
        _general_print(head, color, msg, sys.stderr, level)

def _split_msg(head, msg):
    if isinstance(msg, list):
        msg = '\n'.join(map(str, msg))

    if msg.startswith('\n'):
        # means print \n at first
        msg = msg.lstrip()
        head = '\n' + head

    elif msg.startswith('\r'):
        # means print \r at first
        msg = msg.lstrip()
        head = '\r' + head

    match = PREFIX_RE.match(msg)
    if match:
        head += ' <%s>' % match.group(1)
        msg = match.group(2)

    return head, msg

def get_loglevel():
    return (k for k, v in LOG_LEVELS.items() if v == LOG_LEVEL).next()

def set_loglevel(level):
    global LOG_LEVEL
    if level not in LOG_LEVELS:
        # no effect
        return

    LOG_LEVEL = LOG_LEVELS[level]

def set_interactive(mode=True):
    global INTERACTIVE
    if mode:
        INTERACTIVE = True
    else:
        INTERACTIVE = False

def log(msg=''):
    # log msg to LOG_CONTENT then save to logfile
    global LOG_CONTENT
    if msg:
        LOG_CONTENT += msg

def raw(msg=''):
    _general_print('', NO_COLOR, msg)

def info(msg):
    head, msg = _split_msg('Info', msg)
    _general_print(head, INFO_COLOR, msg)

def verbose(msg):
    head, msg = _split_msg('Verbose', msg)
    _general_print(head, INFO_COLOR, msg, level='verbose')

def warning(msg):
    head, msg = _split_msg('Warning', msg)
    _color_perror(head, WARN_COLOR, msg)

def debug(msg):
    head, msg = _split_msg('Debug', msg)
    _color_perror(head, ERR_COLOR, msg, level='debug')

def error(msg):
    head, msg = _split_msg('Error', msg)
    _color_perror(head, ERR_COLOR, msg)
    sys.exit(1)

def ask(msg, default=True):
    _general_print('\rQ', ASK_COLOR, '')
    try:
        if default:
            msg += '(Y/n) '
        else:
            msg += '(y/N) '
        if INTERACTIVE:
            while True:
                repl = raw_input(msg)
                if repl.lower() == 'y':
                    return True
                elif repl.lower() == 'n':
                    return False
                elif not repl.strip():
                    # <Enter>
                    return default

                # else loop
        else:
            if default:
                msg += ' Y'
            else:
                msg += ' N'
            _general_print('', NO_COLOR, msg)

            return default
    except KeyboardInterrupt:
        sys.stdout.write('\n')
        sys.exit(2)

def choice(msg, choices, default=0):
    if default >= len(choices):
        return None
    _general_print('\rQ', ASK_COLOR, '')
    try:
        msg += " [%s] " % '/'.join(choices)
        if INTERACTIVE:
            while True:
                repl = raw_input(msg)
                if repl in choices:
                    return repl
                elif not repl.strip():
                    return choices[default]
        else:
            msg += choices[default]
            _general_print('', NO_COLOR, msg)

            return choices[default]
    except KeyboardInterrupt:
        sys.stdout.write('\n')
        sys.exit(2)

def pause(msg=None):
    if INTERACTIVE:
        _general_print('\rQ', ASK_COLOR, '')
        if msg is None:
            msg = 'press <ENTER> to continue ...'
        raw_input(msg)

def set_logfile(fpath):
    global LOG_FILE_FP

    def _savelogf():
        if LOG_FILE_FP:
            with open(LOG_FILE_FP, 'w') as log:
                log.write(LOG_CONTENT)

    if LOG_FILE_FP is not None:
        warning('duplicate log file configuration')

    LOG_FILE_FP = fpath

    import atexit
    atexit.register(_savelogf)

def enable_logstderr(fpath):
    global CATCHERR_BUFFILE_FD
    global CATCHERR_BUFFILE_PATH
    global CATCHERR_SAVED_2

    if os.path.exists(fpath):
        os.remove(fpath)
    CATCHERR_BUFFILE_PATH = fpath
    CATCHERR_BUFFILE_FD = os.open(CATCHERR_BUFFILE_PATH, os.O_RDWR|os.O_CREAT)
    CATCHERR_SAVED_2 = os.dup(2)
    os.dup2(CATCHERR_BUFFILE_FD, 2)

def disable_logstderr():
    global CATCHERR_BUFFILE_FD
    global CATCHERR_BUFFILE_PATH
    global CATCHERR_SAVED_2

    raw(msg=None) # flush message buffer and print it.
    os.dup2(CATCHERR_SAVED_2, 2)
    os.close(CATCHERR_SAVED_2)
    os.close(CATCHERR_BUFFILE_FD)
    os.unlink(CATCHERR_BUFFILE_PATH)
    CATCHERR_BUFFILE_FD = -1
    CATCHERR_BUFFILE_PATH = None
    CATCHERR_SAVED_2 = -1