wfb-ng/wfb_ng/__init__.py

import sys
import os
import queue
import threading
import atexit

from twisted.internet import utils, reactor
from logging import currentframe
from twisted.python import log
from twisted.python.logfile import LogFile

__orig_msg = log.msg
_srcfile = os.path.splitext(os.path.normcase(__file__))[0] + '.py'

# Returns escape codes from format codes
esc = lambda *x: '\033[' + ';'.join(x) + 'm'

# The initial list of escape codes
escape_codes = {
    'reset': esc('39', '49', '0'),
    'bold': esc('01'),
}

# The color names
colors = [
    'black',
    'red',
    'green',
    'yellow',
    'blue',
    'purple',
    'cyan',
    'white'
]

# Create foreground and background colors...
for lcode, lname in [('3', ''), ('4', 'bg_')]:
    # ...with the list of colors...
    for code, name in enumerate(colors):
        code = str(code)
        # ...and both normal and bold versions of each color
        escape_codes[lname + name] = esc(lcode + code)
        escape_codes[lname + "bold_" + name] = esc(lcode + code, "01")


def color_str(arg, c, bold=False):
    return '%s%s%s%s' % (escape_codes[c], escape_codes['bold'] if bold else '', arg, escape_codes['reset'])


class ConsoleObserver(object):
    def emit(self, eventDict):
        print('[%s] %s' % (eventDict['system'], eventDict['log_text']))
        sys.stdout.flush()


def __findCaller():
    """
    Find the stack frame of the caller so that we can note the source
    file name, line number and function name.
    """
    f = currentframe()

    #On some versions of IronPython, currentframe() returns None if
    #IronPython isn't run with -X:Frames.

    if f is not None:
        f = f.f_back

    rv = "(unknown file)", 0, "(unknown class)", "(unknown function)"

    while hasattr(f, "f_code"):
        co = f.f_code
        filename = os.path.normcase(co.co_filename)
        if filename == _srcfile:
            f = f.f_back
            continue

        try:
            if 'self' in f.f_locals:
                klass = f.f_locals['self'].__class__.__name__
            else:
                klass = ''
        except:
            klass = '<undef>'

        rv = (co.co_filename, f.f_lineno, klass, co.co_name)
        break
    return rv


class LogLevel(object):
    """Log levels"""
    DEBUG = 1
    INFO = 2
    NOTICE = 3
    WARNING = 4
    ERROR = 5
    ALERT = 6
    FATAL = 7


log_level_map = { LogLevel.DEBUG : 'debug',
                  LogLevel.INFO : 'info',
                  LogLevel.NOTICE : 'notice',
                  LogLevel.WARNING : 'warning',
                  LogLevel.ERROR : 'error',
                  LogLevel.ALERT : 'alert',
                  LogLevel.FATAL : 'fatal_error'}


def _log_msg(*args, **kwargs):
    def _stub():
        return __findCaller()

    level = kwargs.get('level', None)

    if level not in set(LogLevel.__dict__.values()):
        level = LogLevel.ERROR if kwargs.get('isError') else LogLevel.INFO
        kwargs['level'] = level

    error = (level >= LogLevel.ERROR)
    kwargs['isError'] = 1 if error else 0

    if 'why' in kwargs:  # handle call from log.err
        path, line, klass, func = __findCaller()
    else:
        path, line, klass, func = _stub()

    filename = os.path.basename(path)
    module = os.path.splitext(filename)[0]

    tmp = [color_str(module, 'red' if error else 'blue')]

    if klass:
        tmp.append(color_str(klass, 'red' if error else 'blue'))

    tmp.append(color_str(func, 'red' if error else 'green'))
    kwargs['system'] = '%s #%s' % ('.'.join(tmp), log_level_map[level])

    return __orig_msg(*args, **kwargs)


class ExecError(Exception):
    pass


def call_and_check_rc(cmd, *args, **kwargs):
    def _check_rc(_args):
        (stdout, stderr, rc) = _args
        if rc != 0:
            err = ExecError('RC %d: %s %s' % (rc, cmd, ' '.join(args)))
            err.stdout = stdout.strip()
            err.stderr = stderr.strip()
            raise err

        log.msg('# %s' % (' '.join((cmd,) + args),))

        if stdout and kwargs.get('log_stdout', True):
            log.msg(stdout)

        return stdout

    def _got_signal(f):
        f.trap(tuple)
        stdout, stderr, signum = f.value
        err = ExecError('Got signal %d: %s %s' % (signum, cmd, ' '.join(args)))
        err.stdout = stdout.strip()
        err.stderr = stderr.strip()
        raise err

    return utils.getProcessOutputAndValue(cmd, args, env=os.environ).addCallbacks(_check_rc, _got_signal)


class ErrorSafeLogFile(object):
    stderr = sys.stderr
    log_cls = LogFile
    log_max = 1000
    binary = False
    twisted_logger = True
    always_flush = True

    def __init__(self, *args, **kwargs):
        cleanup_at_exit = kwargs.pop('cleanup_at_exit', True)
        self.logfile = None
        self.args = args
        self.kwargs = kwargs

        try:
            self.logfile = self.log_cls(*self.args, **self.kwargs)
        except Exception as v:
            if self.twisted_logger:
                print('Unable to open log file: %s' % (v,), file=self.stderr)
            else:
                log.err(v, 'Unable to open log file: %s(%s, %s)' % (self.log_cls, self.args, self.kwargs), isError=True)

        self.need_stop = threading.Event()
        self.lock = threading.RLock()
        self.log_queue_overflow = 0
        self.log_queue = queue.Queue(self.log_max)
        self.thread = threading.Thread(target=self._log_thread_loop, name='logging thread')
        self.thread.daemon = True
        self.thread.start()

        if cleanup_at_exit:
            atexit.register(self._cleanup)

    def _cleanup(self):
        self.log_queue.join()
        self.need_stop.set()
        self.thread.join()

    def _log_thread_loop(self):
        while not self.need_stop.is_set():
            try:
                data = self.log_queue.get(timeout=1)
            except queue.Empty:
                continue

            overflow = 0
            with self.lock:
                if self.log_queue_overflow:
                    overflow = self.log_queue_overflow
                    self.log_queue_overflow = 0

            if overflow and not self.binary:
                self._write_and_flush('--- Dropped %d log items due to log queue overflow\n' % (overflow,))

            self._write_and_flush(data)
            self.log_queue.task_done()

    def _write_and_flush(self, data):
        try:
            if self.logfile is None:
                self.logfile = self.log_cls(*self.args, **self.kwargs)

            self.logfile.write(data)
            if self.always_flush:
                self.logfile.flush()
        except Exception as v:
            if self.twisted_logger:
                # Don't use logger due to infinite loop
                print('Unable to write to log file: %s' % (v,), file=self.stderr)
            else:
                reactor.callFromThread(log.err, v, 'Unable to write to: %s(%s, %s)' % (self.log_cls, self.args, self.kwargs), isError=True)

            if self.logfile is not None:
                self.logfile.close()
                self.logfile = None

    def write(self, data):
        try:
            self.log_queue.put_nowait(data)
        except queue.Full:
            with self.lock:
                self.log_queue_overflow += 1

    def flush(self):
        pass