Forked from http://paste.ubuntu.com/6124790/

color.py

# -*- encoding: utf-8 -*-

import sys
import time
import traceback

class Color:
    BLACK       = 0
    RED         = 1
    GREEN       = 2
    YELLOW      = 3
    BLUE        = 4
    MAGENTA     = 5
    CYAN        = 6
    WHITE       = 7
    LBLACK      = 8
    LRED        = 9
    LGREEN      = 10
    LYELLOW     = 11
    LBLUE       = 12
    LMAGENTA    = 13
    LCYAN       = 14
    LWHITE      = 15

class Level:
    DEBUG   = 0
    INFO    = 1
    WARN    = 2
    NOTICE  = 3
    ERROR   = 4
    FATAL   = 5 # This log level will trigger a backtrace, and terminate thread with sys.exit(1)
    TRACE   = 6 # This is the level for backtrace only.

class _logger:

    _modeStart  = '\x1B['


    _colors = {
        Level.DEBUG     : Color.LGREEN,
        Level.INFO      : Color.LBLUE,
        Level.NOTICE    : Color.LCYAN,
        Level.WARN      : Color.RED,
        Level.ERROR     : Color.LRED,
        Level.FATAL     : Color.LRED,
        Level.TRACE     : Color.GREEN,
    }

    _levelStr = {
        Level.DEBUG     : 'DEBUG',
        Level.INFO      : 'INFO',
        Level.NOTICE    : 'NOTICE',
        Level.WARN      : 'WARN',
        Level.ERROR     : 'ERROR',
        Level.FATAL     : 'FATAL',
        Level.TRACE     : 'TRACE',
    }

    def __init__(self, f = None, colorize = True, level = None, log_format = None):
        '''
            log_format supports:
            time        -> Time in local timezone. (like what `data` outputs)
            clock       -> Timestamp delta from first call to import logger.
            level       -> Log level
            message     -> What you want to log
        '''
        self._initTime = time.time()
        if level is not None:
            self._level = level
        else:
            self._level = Level.DEBUG
        if f is not None:
            self.fd = f
        else:
            self.fd = sys.stdout
        if log_format is None:
            self._format = '[{clock}][{level}] {message}'
        else:
            self._format = log_format
        self.colorize = colorize and self.fd.isatty()
        self._log(Level.INFO, 'Logger initialized at %s' % time.asctime())

    def _setMode(self, *modes):
        if not self.colorize:
            return
        self.fd.write(self._modeStart)
        self.fd.write(';'.join(map(str, modes)))
        self.fd.write('m')

    def _resetMode(self):
        self._setMode(0)

    def _setColor(self, color):
        if type(color) is not int:
            raise ValueError('Color code should be an int.')
        if color < 0 or color > 15:
            raise ValueError('Non valid color code.')
        if color & 8: # Light colors
            self._setMode(1, (color & 7) + 30)
            return
        self._setMode(color + 30)

    def _setBgColor(self, color):
        if type(color) is not int:
            raise ValueError('Color code should be an int.')
        if color < 0 or color > 15:
            raise ValueError('Non valid color code.')
        if color & 8: # Light colors
            self._setMode(1, (color & 7) + 40)
            return
        self._setMode(color + 40)

    def _log(self, level, message, _traceback = None):
        if level not in self._colors:
            raise ValueError('Non valid level code.')
        if level < self._level:
            return
        self._setColor(self._colors[level])
        _t = time.time()
        _ms = int((_t - int(_t)) * 1000)
        _l = time.localtime()[:6] + (_ms, )
        _time = '%04d-%02d-%02d %02d:%02d:%02d.%03d' % _l
        _clock = '%11.3f' % (_t - self._initTime)
        _level = '%6s' % self._levelStr[level]
        _message = message
        finalMsg = self._format.format(time = _time, clock = _clock, level = _level, message = _message)
        self.fd.write(finalMsg)
        if not finalMsg.endswith('\n'):
            self.fd.write('\n')
        self._resetMode()
        if level == Level.FATAL:
            self._backTrace(_traceback)
            sys.exit(1)

    def _backTrace(self, _frames = None):
        # @TODO: Print backtrace with logger.
        if _frames is None:
            _frames = traceback.extract_stack()[:-3][::-1]
        else:
            _frames = traceback.extract_tb(_frames)
        _depth = 0
        self._log(Level.TRACE, 'CallStack BackTrace:')
        for _frame in _frames:
            _code = _frame[3]
            if _code is None:
                _code = ''
            self._log(Level.TRACE, 'Trace #%d: %s' % (_depth, _code))
            self._log(Level.TRACE, '  Module   : ' + _frame[2])
            self._log(Level.TRACE, '  Location : ' + _frame[0] + (':%d' % _frame[1]))
            _depth += 1
        pass

    def _exceptHandler(self, _type, value, traceback):
        _msg = ''
        if hasattr(value, 'message') and value.message:
            _msg = value.message
        elif hasattr(value, 'args') and value.args:
            _msg = ', '.join(map(str, value.args))
        elif hasattr(value, '__str__'):
            _msg = value.__str__()
        self._log(Level.FATAL, 'Unexpected exception: [%s] %s' % (_type.__name__, _msg), traceback)

# Auto init

_instance = _logger()

def Debug(message):
    _instance._log(Level.DEBUG, message)

def Info(message):
    _instance._log(Level.INFO, message)

def Notice(message):
    _instance._log(Level.NOTICE, message)

def Warn(message):
    _instance._log(Level.WARN, message)

def Error(message):
    _instance._log(Level.ERROR, message)

def Fatal(message):
    _instance._log(Level.FATAL, message)

def _exceptionHandler(_type, _value, _traceback):
    _instance._exceptHandler(_type, _value, _traceback)

sys.excepthook = _exceptionHandler