Script to rotate log files in a directory

rotate.py

#!/usr/bin/env python2.7
# coding: utf-8
'''
This script shows an example of rotating log files in a directory.
'''

# Imports #####################################################################
import os
import sys
import tarfile
from datetime import date


# Metadata ####################################################################
__author__ = 'Timothy McFadden'
__creationDate__ = '26-JUL-2017'


# Globals #####################################################################
def rotated_name(file_path, extention=None):
    '''
    Calculate the rotated file's name.

    The function loops until X.##.AAA does not exist.
    '''
    def get_path(dirname, name_part, index, extention):
        '''Return the path to try'''
        if index:
            fname = '{name_part}-{date}-{index:02}{extention}'.format(
                name_part=name_part, date=date.today().strftime('%Y%m%d'),
                index=index, extention=extention
            )
        else:
            fname = '{name_part}-{date}{extention}'.format(
                name_part=name_part, date=date.today().strftime('%Y%m%d'),
                extention=extention
            )

        return os.path.join(dirname, fname)

    # Split apart the path so we can put it back together with the unique
    # identifier.
    dirname = os.path.dirname(file_path)
    fname = os.path.basename(file_path)

    if extention is None:
        name_part, extention = os.path.splitext(fname)
    else:
        name_part = fname.split(extention)[0]

    # Search for a filename that doesn't exist
    index = 0
    new_path = get_path(dirname, name_part, index, extention)
    while os.path.exists(new_path) and (index < 1):
        index += 1
        new_path = get_path(dirname, name_part, index, extention)

    if os.path.exists(new_path):
        raise Exception("Could not calculate a non-existent path for [%s]" % file_path)

    return new_path


def rotate_logs(log_dir, filter=lambda x: x.endswith('.log'), truncate=False):
    '''
    Rotates all files in a folder matching the ``filter``.

    The rotation means:

        *   Finding a new filname for the file
        *   Creating a zipped tarball with the new filename
        *   Possibly truncating the file (see ``truncate`` parameter)

    WARNING: Any processes that have the logs open should be stopped before
        rotation.

    :param str log_dir: The directory holding the logs
    :param code filter: A function used to filter the logs for rotation
    :param bool truncate: Whether or not to truncate the log after rotation
    '''
    for fname in [x for x in os.listdir(log_dir) if filter(x)]:
        log_path = os.path.join(log_dir, fname)
        archive = rotated_name(log_path + '.tgz')
        with tarfile.open(archive, "w:gz") as tar:
            tar.add(log_path, arcname=fname)

        # Truncate the files
        if truncate:
            with open(log_path, 'wb') as fh:
                fh.truncate()


def main():
    rotate_logs(sys.argv[1])


if __name__ == '__main__':
    if len(sys.argv) != 2:
        print "USAGE: rotate.py <log folder>"
        sys.exit(1)

    main()