test_stopping_mp_dummy_pool.py

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


import signal
from time import sleep
from multiprocessing import TimeoutError
from multiprocessing.dummy import Pool


def work(i):
    print('Working on {} ...'.format(i))
    sleep(1)
    return i


def complicated_main():
    '''Work on both Python 2 and 3.

    Python 2's main thread will wait non-daemon threads, and lock can't be
    interrupt by signal.
    '''

    print('Running complicated main ...')

    pool = Pool(4)
    # hack: make the flag writable in the handler
    pool.ctrl_c_happens = False

    # stop running and waiting if Ctrl-C

    def end_pool(signum, frame):
        print('Ctrl-C happens!')
        pool.terminate()
        pool.ctrl_c_happens = True

    signal.signal(signal.SIGINT, end_pool)

    # run and wait tasks

    # mapr: map result
    mapr = pool.map_async(work, range(8))
    retvals = []

    while True:

        try:
            print('Are tasks finished?')
            retvals = mapr.get(timeout=1)
        except TimeoutError:
            pass
        else:
            break

        print('Does Ctrl-C happen?')
        if pool.ctrl_c_happens:
            break

    print(retvals)


def simple_main():
    '''Work on Python 3.2 up only.

    Python 3.2 up's main thread will also wait non-daemon threads, but lock
    can be interrupt by signal.
    '''

    print('Running simple main ...')

    pool = Pool(4)
    try:
        print(pool.map(work, range(8)))
    except KeyboardInterrupt:
        print([])


if __name__ == '__main__':

    import sys

    if len(sys.argv) >= 2 and sys.argv[1] == 'simple':
        simple_main()
    else:
        complicated_main()