import logging

from pgtoolkit import highlight

class ColoredFormatter(logging.Formatter):

    def __init__(self, fmt, datefmt):

        logging.Formatter.__init__(self, fmt, datefmt)

    def format(self, record):

        color = ()

        if record.levelno == 21:

            color = (2, None)

        if record.levelno == logging.ERROR and not record.exc_info is None:

            color = (11, 1)

        if color:

            origmsg = record.msg

            f,b = color

            record.msg = highlight(1, f, b) + record.msg + highlight(0)

            res = logging.Formatter.format(self, record)

            record.msg = origmsg

        else:

            res = logging.Formatter.format(self, record)

        return res