Splits CSV files by line

csvsplitter.py

import csv
import sys

if __name__ == '__main__':
    if len(sys.argv) != 3:
        sys.exit("%s <csv-file> <lines>" % sys.argv[0])
    infile, total = sys.argv[1], int(sys.argv[2])
    basename = infile.split('.csv')[0]
    output_count = 0
    output_handler, file_handler = None, None
    line_count = 0
    with open(infile, 'rU') as csvinput:
        reader = csv.reader(csvinput)
        header = reader.next()
        for row in reader:
            if line_count >= total:
                file_handler.close()
                output_handler, file_handler = None, None
                line_count = 0
            if output_handler is None:
                output_count += 1
                filename = "%s_%03d.csv" % (basename, output_count)
                file_handler = open(filename, 'w')
                print "writing file:", filename
                output_handler = csv.writer(file_handler)
                output_handler.writerow(header)
            output_handler.writerow(row)
            line_count += 1
    if output_handler is not None:
        file_handler.close()