General utilities

utils.py

#!/usr/bin/env python

"""
Script including useful commands

Includes:
    Flatten folders -- flattens folders into a single folder (can be called recursively)
    Compare         -- compares two strings/files and prints out results
"""

import os
import shutil
import difflib
from argparse import ArgumentParser


def read_file(filename):
    with open(filename) as f:
        return f.read()


def flatten_folders(args):
    """
    Gets the files and folders that need to be processed

    Args:
        :argparse.Namespace args: parsed arguments
    """
    for path in args.input:
        if not os.path.isdir(path):
            print 'Skipping {0} as it is not a directory...'.format(path)
            continue

        path_num_dirs = len(path.split(os.sep))

        # if recursive, use os.walk
        # otherwise, use os.listdir
        for root, dirs, files in os.walk(path):
            num_dirs_diff = len(root.split(os.sep)) - path_num_dirs
            # check to make sure we're withing the number of levels we want to be in
            if root == path:
                continue
            if args.levels and num_dirs_diff > args.levels:
                continue

            print root, num_dirs_diff
            delete_dir = True
            for filename in files:
                new_file = os.path.join(path, filename)
                if os.path.exists(new_file):
                    delete_dir = False
                    print 'Skipping {0} because file with same name exists in {1}'.format(new_file, path)
                    continue
                print 'Moving {0} from {1} to {2}'.format(filename, root, path)
                shutil.move(os.path.join(root, filename), path)

            if delete_dir:
                print 'Removing {0}'.format(root)
                shutil.rmtree(root)
            else:
                print 'Not removing {0} because there\'s still at least one file in it'.format(root)


def compare(args):
    """
    Compare two files or strings

    Args:
        :argparse.Namespace args: parsed arguments
    """
    original = read_file(args.input[0]) if os.path.isfile(args.input[0]) else args.input[0]
    current = read_file(args.input[1]) if os.path.isfile(args.input[1]) else args.input[1]

    diffs_printed = False
    if args.diff:
        for line in difflib.Differ().compare(original.splitlines(True), current.splitlines(True)):
            if 'whole' in args.diff:
                print line
                diffs_printed = True
                continue
            if 'partial' in args.diff and not line.startswith('  '):
                diffs_printed = True
                print line

    if diffs_printed:
        print '\n'
    print 'Percent Match: {0:.2f}%'.format(difflib.SequenceMatcher(None, original, current).ratio() * 100)


def get_arguments():
    """
    Get input arguments

    Returns:
        argparse.Namespace: parsed arguments
    """
    parser = ArgumentParser(description='General utilities')
    subparsers = parser.add_subparsers(help='commands', dest='command')

    # Flatten folders parsers
    flatten_folder_parser = subparsers.add_parser('flatten', help='Flatten folders into one folder')
    flatten_folder_parser.add_argument(dest='input', help='Input directory/directories', nargs='+')
    flatten_folder_parser.add_argument('-l', '--levels', dest='levels', type=int, help='Number of levels to flatten')

    # Compare parser
    compare_parser = subparsers.add_parser('compare', help='Compare two files or strings')
    compare_parser.add_argument(dest='input', help='Input files or strings', nargs=2)
    compare_parser.add_argument('-d', '--diff', dest='diff', nargs='?', help='Show diff', const='partial',
                                choices=['whole', 'partial'])


    return parser.parse_args()


FUNC = {'flatten': flatten_folders,
        'compare': compare}


if __name__ == '__main__':
    ARGS = get_arguments()
    FUNC[ARGS.command](ARGS)