Richard Jones' Log: Improved DEBUG logging only when it's needed

Thanks to Kent Johnson for reminding me about contexts (I don't get to play with them much so I sometimes forget :) usage is a lot simpler:

    for i in range(5):
        with DebugCapture(log):
            task(i) 

This is produced by modifying the code from yesterday a little:

'''Capture a logger's DEBUG output for potential later dumping.

Simple use (requires "with" statement in Python 2.5 or later):

    with DebugCapture(logger):
        do_stuff()

If an exception occurs the DEBUG output will be emitted to the log followed
by an exception ERROR message.

Alternatively for Python < 2.5 you may use it more manually:

    handler = DebugCapture(logger)

And at the end of a transaction:

    handler.clear()

If an error occurs:

    handler.emit_debug()

And to uninstall (stop trapping DEBUG messages):

    handler.uninstall()

Copyright 2009 Richard Jones (richard@mechanicalcat.net) and you're free
to use / modify / ignore it as you like.
'''

import logging
import weakref

class DebugCapture(logging.Handler):
    def __init__(self, logger):
        logging.Handler.__init__(self, logging.DEBUG)
        self.records = []
        # insert myself as the handler for the logger
        self.logger = weakref.proxy(logger)
        self.slaves = logger.handlers
        logger.handlers = [self]
    def emit_debug(self):
        # emit stored records to the original logger handler(s)
        for record in self.records:
            for slave in self.slaves:
                slave.handle(record)
    def clear(self):
        self.records = []
    def uninstall(self):
        self.logger.handlers = self.slaves
        self.slaves = self.logger = None
    def emit(self, record):
        self.records.append(record)
    def handle(self, record):
        # if its a DEBUG level record then intercept otherwise
        # pass through to the original logger handler(s)
        if record.levelno == logging.DEBUG:
            return logging.Handler.handle(self, record)
        return sum(slave.handle(record) for slave in self.slaves)
    def __enter__(self):
        pass
    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is not None:
            self.emit_debug()
            self.logger.error('unhandled exception', exc_info=(exc_type,
                exc_value, traceback))
        self.uninstall()
        return True

#
# SAMPLE PROGRAM
#
logging.basicConfig()
log = logging.getLogger()
log.setLevel(logging.DEBUG)

def task(argument):
    log.info('run task (argument=%s)'%argument)
    log.debug('debug info should only appear next to error')
    if argument%2: BROKEN
    log.debug('all done!')

if __name__ == '__main__':
    for i in range(5):
        with DebugCapture(log):
            task(i)