python multiprocessing example

multiprocessing.py

#!/usr/bin/env python3

""" Example multiprocess based python file
"""

import multiprocessing
import argparse
import itertools

ARGP = argparse.ArgumentParser(
    description=__doc__,
    formatter_class=argparse.RawTextHelpFormatter,
)
ARGP.add_argument('--process-count', '-T', type=int,
                  help='''Specific multiprocessor process count, 1 will disable multiprocessing. (Default: cpu count)''')


def foo(bar, baz):
    return bar*baz


def _mp_example(processes, worker, jobs_kwargs):
    """ MP with single process mode
    """
    if processes == 1:
        for job_kwargs in jobs_kwargs:
            yield worker(**job_kwargs)
        return


    # start 4 worker processes
    with multiprocessing.Pool(processes=processes) as pool:
        try:
            job_results = [
                pool.apply_async(worker, kwds=job_kwargs)
                for job_kwargs in jobs_kwargs
            ]

            for job_result in job_results:
                yield job_result.get()
        except multiprocessing.TimeoutError as err:
            print(err)
        except KeyboardInterrupt as err:
            print(err)


def _collect_work():
    for bar_int, baz_int in itertools.product(range(10), repeat=2):
        yield {'bar': bar_int, 'baz': baz_int}


def main(argp=None):
    """ CLI Entry
    """
    if argp is None:
        argp = ARGP.parse_args()  # pragma: no cover


    for result in _mp_example(argp.process_count, foo, _collect_work()):
        print(result)


if __name__ == '__main__':
    main()  # pragma: no cover