LZMA compressed RotatingFileHandler with support for backupCount

lzmarotate.py

"""Example extension of RotatingFileHandler that LZMA compresses the rotated files

Other examples I found didn't actually respect maxBackups, this does

- AG
"""
import codecs
import glob
import logging.handlers
import os
import time
import lzma
from os.path import exists

MAX_LOG_BACKUP = 5
MAX_LOG_SIZE = 2 << 23  # 16MB


class CompressingRotatingFileHandler(logging.handlers.RotatingFileHandler):
    """Rotating log handler that uses an arbitrary compression algorithm after rolling files

    By default, it will use the lzma algorithm and will use the extension ".lzma"
    The behavior is otherwise identical to a standard RotatingFileHandler
    If backupCount > 1, the files will be named as:

    logfile.log.1.<ext>
    logfile.log.2.<ext>
    logfile.log.3.<ext>
    ...

    ... where the larger the number, the older the file. I've seen much simpler versions of this
    that only override rotator() but they don't work when backupCount > 1; without the additional
    prerotate logic, the logfile.log.1.<ext> file will continuously be overwritten and no other
    copies will be kept

    Examples
    --------
    Usage is identical to logging.handlers.RotatingFileHandler:
        >>> import lzma
        >>> file_handler = CompressingRotatingFileHandler(
            'out.log', maxBytes=MAX_LOG_SIZE, backupCount=MAX_LOG_BACKUP,
            compressor_opener=lzma.open compressor_ext='lzma')
        >>> file_handler.setLevel(logging.DEBUG)
        >>> file_formatter = logging.Formatter('%(levelname)s %(funcName)s():%(lineno)d\t%(message)s')
        >>> file_handler.setFormatter(file_formatter)
        >>> logging.getLogger().addHandler(file_handler)

    Notes
    -----
    Change compressor_opener and compressor_ext to use different algorithms as noted in the constructor
    docstring

    - AG
    """
    def __init__(self, *args, compressor_opener=lzma.open, compressor_ext='lzma', **kwargs):
        """Pass optional kwargs compressor_opener and _compressor_ext, or keep lzma as the default

        compressor: callable
            A callable that is called as open(), accepting a filename and a mode. Suitable
            examples include:
                lzma.open
                gzip.open
                zipfile.ZipFile

        zipfile.ZipFile is actually a class, but it is also a callable that behaves as open() does
        """
        self._compressor_opener = compressor_opener
        self._compressor_ext = compressor_ext
        if not self._compressor_ext.startswith('.'):
            self._compressor_ext = '.' + self._compressor_ext
        super(CompressingRotatingFileHandler, self).__init__(*args, **kwargs)
        self.rotator = self._rotator

    def _prerotate(self) -> None:
        """This is to move the backed up and compressed files in the case that backupCount > 1

        Working backwards from the highest value, this moves:
            <baseFilename>.<n>.<ext> -> <baseFilename>.<n+1>.<ext>
            ...

        There is similar logic inside doRollover() but it's simpler to reproduce it here with a slight
        modification (the compressor extension)
        """
        if self.backupCount > 0:  # noqa
            for i in range(self.backupCount - 1, 0, -1):  # noqa
                sfn = self.rotation_filename(f'{self.baseFilename}.{i}{self._compressor_ext}')  # noqa
                dfn = self.rotation_filename(f'{self.baseFilename}.{i + 1}{self._compressor_ext}')  # noqa
                if os.path.exists(sfn):
                    if os.path.exists(dfn):
                        os.remove(dfn)
                    os.rename(sfn, dfn)
            dfn = self.rotation_filename(self.baseFilename + '.1')  # noqa
            if os.path.exists(dfn):
                os.remove(dfn)

    def _rotator(self, src: str, dst: str) -> None:
        """This overrides the rotator method to actually performs the compression"""
        # Move any <n>.<extension> to <n+1>.<extension>
        self._prerotate()
        if not exists(src):
            # This could have issues with race conditions, but otherwise I don't think it's a thing
            # Will see if it ever fires and fix it later - AG
            raise RuntimeError('Is this a thing? Or is this only called with files that exist?')
        os.rename(src, dst)
        outfile = dst + self._compressor_ext
        with open(dst, mode='rb') as f_in, self._compressor_opener(outfile, mode='wb') as f_out:
            f_out.writelines(f_in)
        os.remove(dst)


class CompressingTimedRotatingFileHandler(logging.handlers.TimedRotatingFileHandler):
    """Timed rotator that also does compression; mostly stolen from StackOverflow"""
    def __init__(self, *args, compressor_ext='lzma', compressor_opener=lzma.open, **kwargs):
        self._compressor_opener = compressor_opener
        self._compressor_ext = compressor_ext
        if not self._compressor_ext.startswith('.'):
            self._compressor_ext = '.' + self._compressor_ext
        super(CompressingTimedRotatingFileHandler, self).__init__(*args, **kwargs)
        self.rotator = self._rotator

    def _rotator(self, source: str, dest: str) -> None:
        """Implement the log roll, a simple rename"""
        os.rename(source, dest)
        with open(dest, 'rb') as infd, lzma.open(dest + self.ext, mode='wb') as outfd:
            outfd.writelines(infd)
        os.remove(dest)

    @property
    def ext(self):
        return self._compressor_ext

    def _prerotate(self) -> None:
        """This is to move the backed up and compressed files in the case that backupCount > 1

        Working backwards from the highest value, this moves:
            <baseFilename>.<n>.<ext> -> <baseFilename>.<n+1>.<ext>
            ...

        There is similar logic inside doRollover() but it's simpler to reproduce it here with a slight
        modification (the compressor extension)
        """
        if self.backupCount > 0:  # noqa
            for i in range(self.backupCount - 1, 0, -1):  # noqa
                sfn = self.rotation_filename(f'{self.baseFilename}.{i}{self._compressor_ext}')  # noqa
                dfn = self.rotation_filename(f'{self.baseFilename}.{i + 1}{self._compressor_ext}')  # noqa
                if os.path.exists(sfn):
                    if os.path.exists(dfn):
                        os.remove(dfn)
                    os.rename(sfn, dfn)
            dfn = self.rotation_filename(self.baseFilename + '.1')  # noqa
            if os.path.exists(dfn):
                os.remove(dfn)

    def doRollover(self):
        """
        do a rollover; in this case, a date/time stamp is appended to the filename
        when the rollover happens.  However, you want the file to be named for the
        start of the interval, not the current time.  If there is a backup count,
        then we have to get a list of matching filenames, sort them and remove
        the one with the oldest suffix.
        """
        if hasattr(self.stream, 'close'):
            self.stream.close()
        else:
            raise RuntimeError('Invalid stream!')
        # get the time that this sequence started at and make it a TimeTuple
        t = self.rolloverAt - self.interval  # noqa
        time_tuple = time.localtime(t)
        dfn = self.baseFilename + "." + time.strftime(self.suffix, time_tuple)  # noqa
        if os.path.exists(dfn):
            os.remove(dfn)
        os.rename(self.baseFilename, dfn)  # noqa
        if self.backupCount > 0:  # noqa
            # find the oldest log file and delete it
            # Probably should fix this hard-coded glob string ...
            s = glob.glob(self.baseFilename + ".20*")  # noqa
            if len(s) > self.backupCount:  # noqa
                s.sort()
                os.remove(s[0])
        if self.encoding:  # noqa
            self.stream = codecs.open(self.baseFilename, mode='w', encoding=self.encoding)  # noqa
        else:
            self.stream = open(self.baseFilename, mode='w')  # noqa
        self.rolloverAt = self.rolloverAt + self.interval  # noqa
        full_dfn = dfn + self.ext
        if os.path.exists(full_dfn):
            os.remove(full_dfn)
        file = self._compressor_opener(full_dfn, mode="wb")
        file.write(dfn)
        file.close()
        os.remove(dfn)

    def __del__(self):
        if hasattr(self, 'stream') and hasattr(self.stream, 'close'):
            try:
                self.stream.close()
            except Exception:  # noqa, pylint: disable=broad-except
                pass