Logger capturing utility

capturelog.py

class CaptureLog:
    """Inside the context of this object, all log from the given logger is redirected to a temporary file.
    Within the context, the file can be read. After the context exits, the file is deleted.

    :param logger: A Logger instance or a string which is the path to the logger.
    :param level: Minimum level to capture logs.
    :param log_format: Format string used to capture logs from this logger.

    Example of usage:

        with capture_log_to_file('logger.name', level=logging.ERROR) as captured:
            ... do things that log to logger.name ...
            if captured.log_size():
                print('There were errors:')
                print(captured.get_content())
    """

    def __init__(self, logger, level=logging.INFO, log_format=None):
        if isinstance(logger, str):
            logger = logging.getLogger(logger)

        self.logger = logger
        self.level = level
        self.log_format = log_format or '%(asctime)s %(message)s'
        self.encoding = 'utf-8'
        self.tmplogfile = None
        self.filename = None
        self.handler = None

    def __repr__(self):
        return f'{self.__class__.__name__}({self.logger.name})'

    def __enter__(self):
        self.tmplogfile = tempfile.NamedTemporaryFile(
            prefix=self.logger.name + '.', suffix='.log', mode='w+', encoding=self.encoding)
        self.handler = logging.StreamHandler(self.tmplogfile)
        self.handler.setFormatter(logging.Formatter(self.log_format))
        self.handler.setLevel(self.level)

        self.logger.addHandler(self.handler)
        return self

    def __exit__(self, exception_type, exception_value, traceback):
        self.logger.removeHandler(self.handler)
        self.handler = None
        self.tmplogfile.close()
        self.tmplogfile = None

    def _check_in_context(self):
        assert self.tmplogfile is not None, "This function can only be called in a CaptureLog context."

    def log_size(self):
        self._check_in_context()
        return self.tmplogfile.tell()

    def get_content(self):
        """Returns all log content in a single string."""
        self._check_in_context()
        with open(self.tmplogfile.name, encoding=self.encoding) as fileobj:
            return fileobj.read()

    def log_sample(self, lines_before, lines_after, delimiter='...'):
        """Returns up to `lines_before + lines_after` lines from the log file.
        If the log has more lines than that, only the first `lines_before` and last `lines_after` are obtained,
        and `delimiter` is used to concatenate them.
        """
        self._check_in_context()
        all_lines = [line for line in open(self.tmplogfile.name, encoding=self.encoding)]
        if len(all_lines) > lines_before + lines_after:
            all_lines = all_lines[:lines_before] + [delimiter + '\n'] + all_lines[-lines_after:]
        return ''.join(all_lines)