multiwrite.py

#!/usr/bin/eng python
#-*- coding: utf-8 -*-
'''
'''

from multiprocessing import Pool
import sys
import os

def task(path):
    try:
        with open(path, 'w') as f:
            f.write('%s: All Set\n' % os.getpid())
    except:
        return False
    return True

def main(args):
    '''
    Writes to a bunch of files in parallel
    '''

    # If the directory doesn't exist create it
    if not os.path.exists(args.output):
        os.makedirs(args.output)

    # Create a pool of workers
    pool = Pool(processes=args.processes)

    # Create a list of the files to write to
    # Each worker will need input
    paths = []
    for i in xrange(args.files):
        path = os.path.join(args.output, str(i) + '.txt')
        paths.append(path)

    # Hand out the jobs to each worker
    results = [pool.apply_async(task, (path,)) for path in paths]

    # Wait on all the children
    for result in results:
        result.get(timeout=10)

    return 0

if __name__ == '__main__':
    from argparse import ArgumentParser

    parser = ArgumentParser(description=main.__doc__)
    parser.add_argument('-n', '--processes', default=4, type=int, help='Number of processes to use')
    parser.add_argument('-o', '--output', default='/tmp/jobs', help='Where to write the files')
    parser.add_argument('files', type=int, help='How many files to write')

    args = parser.parse_args()

    sys.exit(main(args))