split.py

#!/usr/bin/python
import argparse

def generate_argparser():
    parser = argparse.ArgumentParser()
    parser.add_argument('filename', type=str,
                        help="the name of the file to be split")
    parser.add_argument('criteria', nargs='+', type=str,
                        help='a list of space delinated strings to split the file')
    parser.add_argument('-b', '--blank', help='Split file on blank lines', action='store_true')
    return parser


def write_file(name, count, chunk, ext):
    # Generates a new file name, and writes the contents of chunk into a new file
    newfile = name + str(count)
    count += 1
    if ext:
        newfile += '.' + ext
    with open(newfile, 'w+') as nf:
        nf.writelines(chunk)

def split(fname, criteria = [' ']):

    if criteria == []:
        return

    count = 0

    name = fname.split('.')
    ext = None
    if len(name) > 1:
        ext = name[1]
        name = name[0]

    max_len = 1
    for c in criteria:
        max_len = len(c) if len(c) > max_len else max_len

    with open(fname) as infile:
        chunk = ''
        for l in infile.readlines():
            if l[0:max_len] in criteria:
                write_file(name, count, chunk, ext)
                chunk = l
                count += 1
            else:
                chunk += l
    # Write the last chunk of file
    write_file(name, count, chunk, ext)

if __name__ == '__main__':

    parser = generate_argparser()
    args = parser.parse_args()
    if args.blank:
        args.criteria.append('\n')
    split(args.filename, args.criteria)