HoldingLogger

import logging
from contextlib import contextmanager


class HoldingHandler(logging.Handler):
    """ Does never flush the logs if not explicitely asked.

    Keep them all in memory
    """

    def __init__(self, level=logging.NOTSET, target=None):
        """
        Initialize the handler with the buffer size.
        """
        super(HoldingHandler, self).__init__(level=level)
        self.buffer = []
        self.target = target

    def emit(self, record):
        """
        Emit a record.

        Append the record. If shouldFlush() tells us to, call flush() to process
        the buffer.
        """
        self.buffer.append(record)

    def setTarget(self, target):
        """
        Set the target handler for this handler.
        """
        self.target = target

    def flush(self):
        """
        For a HoldingHandler, flushing means just sending the buffered
        records to the target, if there is one. Override if you want
        different behaviour.
        """
        self.acquire()
        try:
            if self.target:
                for record in self.buffer:
                    self.target.handle(record)
                self.buffer = []
        finally:
            self.release()

    def clear(self):
        """
        Empty the buffer. All log lines are discarded.
        """
        self.acquire()
        try:
            self.buffer = []
        finally:
            self.release()

    def close(self):
        """
        Flush, set the target to None and lose the buffer.
        """
        self.acquire()
        try:
            self.target = None
            # need to clear otherwise it gonna be flushed
            self.buffer = []
            super(HoldingHandler, self).close()
        finally:
            self.release()


class HoldingLogger(logging.getLoggerClass()):
    """ A special logger allowing to retain log messages until an exception
    occurs.

    When the context manager ``log_on_error`` is used, all the messages
    will be hold in memory and flushed only if an exception occurs.
    The log messages are discarded at the closing if nothing bad happened.

    This is useful when a - maybe long - workflow is happening and we
    don't want all its details when it goes well. As soon as something
    goes wrong, we'll be able to log a lot of details about the
    workflow.
    """

    @contextmanager
    def flush_on_error(self, target_handler):
        handler = HoldingHandler(target=target_handler)
        current_handlers = self.handlers
        self.handlers = []
        for handler in self.handlers[:]:
            self.removeHandler(handler)
        self.addHandler(handler)
        try:
            yield
        except BaseException:
            handler.flush()
            raise
        finally:
            handler.close()
            for handler in current_handlers:
                self.addHandler(handler)


_origin_logger_class = logging.getLoggerClass()
logging.setLoggerClass(HoldingLogger)
try:
    log = logging.getLogger(__name__)
finally:
    logging.setLoggerClass(_origin_logger_class)
    del _origin_logger_class


log.setLevel(logging.DEBUG)
handler = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
log.addHandler(handler)

log.info('normal log')


with log.flush_on_error(handler):
    log.info('works well 1')
    log.info('works well 2')
    log.info('works well 3')


try:
    with log.flush_on_error(handler):
        log.info('works badly 1')
        log.info('works badly 2')
        log.error('works badly, oups')
        raise Exception('error')
except:
    pass


log.info('normal log')