Task logger classes and decorator

tasklogger

#!/usr/bin/env python
import logging
import re
from collections import (
    OrderedDict,
    deque,
)
from functools import wraps


START_TASK_MSG = 'Starting %s'
START_TASK_TRIGGER_MSG = 'start task %s'
START_TASK_REG = re.compile('start task (?P<task_name>.*)')
END_TASK_MSG = 'Finished %s'
END_TASK_TRIGGER_MSG = 'end task %s'
END_TASK_REG = re.compile('end task (?P<task_name>.*)')



class ColorFormatter(logging.Formatter):
    DEFAULT = '\x1b[0m'
    RED     = '\x1b[31m'
    GREEN   = '\x1b[32m'
    YELLOW  = '\x1b[33m'
    CYAN    = '\x1b[36m'
    CRITICAL = RED
    ERROR    = RED
    WARNING  = YELLOW
    INFO     = GREEN
    DEBUG    = CYAN

    def format(self, record):
        level = record.levelno

        if level >= logging.CRITICAL:
            color = self.CRITICAL
        elif level >= logging.ERROR:
            color = self.ERROR
        elif level >= logging.WARNING:
            color = self.WARNING
        elif level >= logging.INFO:
            color = self.INFO
        elif level >= logging.DEBUG:
            color = self.DEBUG
        else:
            color = self.DEFAULT

        return (
            color
            + super(ColorFormatter, self).format(record)
            + self.DEFAULT
        )


class TaskHandler(logging.StreamHandler):

    def __init__(
        self,
        level=logging.NOTSET,
        formatter=None,
        task_tree_depth=2,
        dump_level=logging.ERROR,
        buffer_size=2000,
        with_color=True,
    ):
        super(TaskHandler, self).__init__(
        )
        self.formatter = formatter
        self.tasks = OrderedDict()
        self.cur_task = None
        self.got_error = False
        self.dump_level = dump_level
        self.buffer_size = buffer_size
        self.task_tree_depth = task_tree_depth
        self.with_color = with_color

    def handle_new_task(self, task_name, record):
        record.msg = START_TASK_MSG % task_name

        self.tasks[task_name] = deque(maxlen=self.buffer_size)
        if self.task_tree_depth >= len(self.tasks):
            super(TaskHandler, self).emit(record)

        self.cur_task = task_name

    def handle_closed_task(self, task_name, record):
        if task_name not in self.tasks:
            return

        record.msg = END_TASK_MSG % task_name

        if self.task_tree_depth >= len(self.tasks):
            super(TaskHandler, self).emit(record)

        to_remove = []
        for cur_task in reversed(self.tasks):
            if self.got_error:
                for record in self.tasks[cur_task]:
                    if record.levelno >= self.level:
                        super(TaskHandler, self).emit(record)

            to_remove.append(cur_task)
            if cur_task == task_name:
                break

        for task in to_remove:
            del self.tasks[task]

        if self.tasks:
            self.cur_task = reversed(self.tasks.keys()).next()
        else:
            self.cur_task = None

    def dump_current_task(self):
        for old_record in self.tasks[self.cur_task]:
            super(TaskHandler, self).emit(old_record)


    def emit(self, record):
        record.task = self.cur_task

        is_start = START_TASK_REG.match(record.msg)
        if is_start:
            self.handle_new_task(is_start.groupdict()['task_name'], record)
            return

        is_end = END_TASK_REG.match(record.msg)
        if is_end:
            self.handle_closed_task(is_end.groupdict()['task_name'], record)
            return

        if record.levelno >= self.dump_level:
            self.got_error = True
            self.dump_current_task()


        if self.got_error:
            if record.levelno >= self.level:
                super(TaskHandler, self).emit(record)
            return

        if self.cur_task and record.levelno >= self.level:
            self.tasks[self.tasks.keys()[-1]].append(record)
            return
        elif self.cur_task:
            return

        if record.levelno >= self.level:
            super(TaskHandler, self).emit(record)


class LogTask(object):
    def __init__(self, task):
        self.task = task

    def __enter__(self):
        logging.info(START_TASK_TRIGGER_MSG % self.task)

    def __exit__(self, *args, **kwargs):
        logging.info(END_TASK_TRIGGER_MSG % self.task)


def log_task(task):
    def decorator(func):

        @wraps(func)
        def wrapper(*args, **kwargs):
            with LogTask(task):
                func(*args, **kwargs)

        return wrapper
    return decorator


@log_task('popo')
def do_popo(msg=None, with_error=False):
    logging.info('Pre popo log')
    if msg and with_error:
        logging.error(msg)
    elif msg:
        logging.info(msg)
    logging.info('Post popo log')


@log_task('popo2')
def do_popo2(msg=None, with_error=False):
    do_popo(msg, with_error)
    logging.info('awonder popo2')




logging.basicConfig(level=logging.INFO)
logging.root.handlers = [
    TaskHandler(
        level=logging.INFO,
        dump_level=logging.ERROR,
        formatter=ColorFormatter(fmt='%(levelname)s::%(msg)s')
    )
]
logging.info('non.task')
do_popo2('nested2!', with_error=True)
logging.info('minimal perlas')