changemon.py

#!/usr/bin/env python
# TODO: timestamps in monitor mode output

DESC="""  changemon.py - Compare or monitor filesystem directory changes

Changemon produces an easy-to-read summary of changes in two modes: In the
first, two directories are compared. In the second mode, a single directory is
monitored for changes as they occur.
"""

NOTES="""
Notes: Changed files are determined by comparing file sizes, and where two
files have the same size, a comparison of their md5 checksums is done (unless
--skip-checksums is supplied). Md5 checks are not done in watch mode.

Flags:

Use flags to select which information to show, The --flags options takes:

a - added files
r - removed
c - changed
s - files in common to both directories
u - files shared by both directories that have not changed

"Watch mode" does not support --flags "su". Ommitting --flags produces the same
result as --flags "acr". The order of flags will be preserved, so --flags "acr"
order the results differently than --flags "car".

Examples:

Compare foo_v1 to foo_v2 to show files that were changed, added and removed:
$changemon.py -c ~/projects/foo_v1 ~/projects/foo_v2

Watch the current directory for changes as they occur
$ changemon.py -w

Watch ~/projects/foo
$ changemon.py -w ~/projects/foo

Author: github.com/nicofanu

"""

import argparse
import hashlib
import os
import sys
import time
import operator
from functools import partial

join, commonprefix = os.path.join, os.path.commonprefix

class monitor(object):
    files_stats = {}

    def __init__(self, criteria, target, interval):
        self.criteria   = criteria
        self.target     = target
        self.interval   = interval
        self.before     = []
        self.after      = []
        self.initial_tree = self.get_tree()

    def get_tree(self):
        tree = list(os.walk(self.target, followlinks=True, onerror=report))
        return collapse(tree)

    def get_stats(self, tree):
        stats = {}
        for f in tree:
            if not f.endswith(os.path.sep):
                fpath = join(self.target, f)
                stats[f] = {'size': os.path.getsize(fpath),
                            'mtime': os.path.getmtime(fpath)}
        return stats

    def start(self):
        stats = monitor.files_stats
        stats['before'] = {}
        stats['after'] = {}
        while True:
            self.before = self.after[:] if self.after else self.initial_tree
            stats['before'] = stats['after'].copy() if stats['after'] else self.get_stats(self.before)

            time.sleep(self.interval)

            self.after = self.get_tree()
            stats['after'] = self.get_stats(self.after)

            compared = comparison(criteria, self.before, self.after)
            pretty_watch(compared)

    def stop(self):
        print
        sys.exit(0)


def parse_args():
    ap = argparse.ArgumentParser(description=DESC,
                                 epilog=NOTES,
                                 formatter_class=argparse.RawTextHelpFormatter)

    mode = ap.add_mutually_exclusive_group(required=True)

    ap.add_argument('-f', '--flags',
                    help='specifies file groups to show, see "flags" below',
                    default='car')
    ap.add_argument('-s', '--skip-checksum',
                    help='don\'t checksum, check modification times instead',
                    action='store_const', const=True, default=False)
    mode.add_argument('-c', '--compare',
                    help='compare two directories',
                    nargs=2, metavar='DIR')
    mode.add_argument('-w', '--watch',
                    help='watch a directory continuously for changes',
                    nargs='?', metavar='DIR', default=os.getcwd())
    ap.add_argument('-i', '--interval',
                    help='update interval in seconds for watch mode',
                    metavar='N', type=int, default=5)
    ap.add_argument('-x', '--cutoff',
                    help='amount of files to list for each group summary in watch mode',
                    metavar='N', type=int, default=3)

    return ap.parse_args()


def memoize(fn):
    def inner(*args, **kwargs):
        key = (args, kwargs)
        first_run = 'cache' not in fn.__dict__
        if first_run:
            result = fn(*args, **kwargs)
            fn.cache = [(key, result)]
        else:
            key_match = lambda tup: key == tup[0]
            match = filter(key_match, fn.cache)
            if match:
                result = match[0][1]
            else:
                result = fn(*args, **kwargs)
                fn.cache.append((key, result))
        return result
    return inner
            

def collapse(tree, strip_root=True):
    """ Takes a list of 3-tuples (as generated by os.walk) and returns a flat list 

    Directories are signified by trailing slashes.
    
        >>> listing = list(os.walk('.')) 
        >>> listing
        [('.', ['.gitignore', 'tags', 'changemon.py'], ['.git']),
        ('.git', ..., ...)]
        >>> collapse(listing)
        ['.gitignore', 'tags', 'changemon.py',
        '.git/', ..., ...]
        >>> collapse(listing, strip_root=False)
        ['./', './.gitignore', './tags', './changemon.py',
        './.git/', ..., ...]
    """
    
    add_sep = lambda s: join(s, '')     # adds path separator... do NOT use os.path.sep!

    genesis = add_sep(tree[0][0])
    
    if strip_root:
        collapsed = []
        add_prefix = lambda root, f: join(root, f).replace(genesis, '', 1)
    else:
        collapsed = [genesis]
        add_prefix = lambda root, f: join(root, f)

    for root, dirs, files in tree:
        files_and_dirs = files + map(add_sep, dirs)
        collapsed.extend([add_prefix(root, f) for f in files_and_dirs])
    
    return collapsed


def collapse_args(strip_root=True):
    """ Decorator to collapse trees as necessary for our comparison functions
    
    This allows us to pass in either lists of 3-tuples from os.walk(), or such
    lists that have already been run through collapse(), to our comparison
    functions. """

    def decorator(fn):
        def inner(before, after):
            _collapse = partial(collapse, strip_root=strip_root)
            if type(before[0]) == tuple:
                before, after = _collapse(before), _collapse(after)
            return fn(before, after)
        return inner
    return decorator


def strip_roots(seq):           return [s.replace(seq[0], '', 1) for s in seq][1:]

########## Comparison functions ################################################

@collapse_args()                       # a - b
def added(before, after):       return [p for p in after if not p in before]

@collapse_args()                       # b - a
def removed(before, after):     return [p for p in before if p not in after]

@collapse_args()
@memoize                               # intersection(b, a)
def common(before, after):      return [p for p in before if p in after]

@collapse_args(strip_root=False)
@memoize
def shared_files(before, after):       # files in common excluding directories
    before, after = strip_roots(before), strip_roots(after)
    return [p for p in common(before, after) if not p.endswith(os.path.sep)] 

@collapse_args(strip_root=False)
@memoize
def changed(before, after):            # filter common_files by differing stats
    _shared_files       = shared_files(before, after)
    root1, root2        = before[0], after[0]
    common_with_roots   = [(join(root1, p), join(root2, p)) for p in _shared_files]
    changed_with_roots  = filter(differing, common_with_roots)
    return strip_roots([root1] + [p[0] for p in changed_with_roots])

@collapse_args(strip_root=False)
def changed_stateful(before, after):
    stats = monitor.files_stats     # functional programming is hard, okay!
    _shared_files = shared_files(before, after)
    _changed = []
    for f in _shared_files:
        if f in stats['before']:
            old = stats['before'][f]
            new = stats['after'][f]
            if old['size'] != new['size'] or old['mtime'] != new['mtime']:
                _changed.append(f)
    return _changed

@collapse_args(strip_root=False)       # common_files - changed_files
def unchanged(before, after):           
    _shared_files       = shared_files(before, after)
    _changed            = changed(before, after)
    return [p for p in _shared_files if not p in _changed]

################################################################################

def comparison(criteria, tree1, tree2):
    """ Return differences between two trees as requested in criteria
    
        >>> criteria
        [('Added', added),      # added is a function
         ('Removed', removed)]
        >>> comparison(tree1, tree2, criteria)
        [('Added', [files...]),
         ('Removed', [files...])]

    """

    results = []
    for label, fn in criteria:
        data = fn(tree1, tree2)
        results += [(label, sorted(data))]
    
    return results


def pretty_compare(results):
    """ Pretty-print for comparison mode """

    def make_group(tup):
        label, seq = tup
        if seq:
            files_string = '\n'.join(seq) + '\n'
        else:
            files_string = ''
        return '{}\n{}\n{}'.format(label, '-' * len(label), files_string)

    output = map(make_group, results)

    print '\n'.join(output)


def pretty_watch(results, cutoff=3):
    """ Pretty-print for watch mode """

    def make_group(tup):
        label, seq = tup
        if len(seq) <= cutoff:
            files_string = ', '.join(seq)
        else:
            files_string = '{} and {} more'.format(', '.join(seq[:cutoff]),
                                                   len(seq[cutoff:]))
        return '{}: {}'.format(label, files_string)

    output = map(make_group, filter(lambda tup: tup[1], results))

    if output: print '\n'.join(output)


def differing(files, skip_checksum=False):
    """ Return True when two files differ - "files" must be a tuple """

    # return file1 != file2
    check = lambda fn, file1, file2: operator.ne(*map(fn, [file1, file2]))

    # if sizes differ, it's true that the files aren't identical
    if check(os.path.getsize, *files):
        return True

    # if sizes don't differ, check md5sums for accuracy
    elif skip_checksum == False:
        return check(checksum, *files)

    # or check modification time instead
    elif skip_checksum == True:
        return check(os.path.getmtime, *files)


def checksum(f):
    with open(f, 'rb') as fh:
        contents = fh.read()
    m = hashlib.md5()
    m.update(contents)
    return m.digest()


def report(err): raise err # Error-reporting function for os.walk() 


if __name__ == '__main__':

    args         = parse_args()

    differing    = partial(differing, skip_checksum=args.skip_checksum)
    pretty_watch = partial(pretty_watch, cutoff=args.cutoff)

    flags        = args.flags.lower()
    initial      = lambda tup: tup[0][0]
    get_flagged  = partial(filter, lambda s: initial(s).lower() in flags)
    sort_by_flag = partial(sorted, key=lambda s: flags.find(initial(s)))

    # comparison mode: compare two directories
    if args.compare:

        table = [('Added',      added),
                 ('Changed',    changed),
                 ('Removed',    removed),
                 ('Shared',     common),
                 ('Unchanged',  unchanged)]

        criteria  = sort_by_flag(get_flagged(table)) 

        make_tree = lambda d: list(os.walk(d, followlinks=True, onerror=report))
        trees     = map(make_tree, args.compare)

        pretty_compare(comparison(criteria, *trees))

    # watch mode: monitor a directory for changes in a continuous loop
    elif args.watch:

        table = [('Added',      added),
                 ('Changed',    changed_stateful),
                 ('Removed',    removed)]

        criteria = sort_by_flag(get_flagged(table)) 
        
        Monitor = monitor(criteria, args.watch, args.interval)

        try:
            Monitor.start()
        except KeyboardInterrupt:
            Monitor.stop()