from multiprocessing.pool import ThreadPool
from multiprocessing.pool import Pool as ProcessPool

def prn(x):
    return "returned(%s)" % x

def exception_throwing_generator(total, when):
    for i in xrange(total):
        if i == when:
            print "exception_throwing_generator(%d, %s) is raising an exception." % (total, when)
            raise Exception("Exception thrown from exception_throwing_generator")
        yield i

def test(pool, total, when):
    print "Testing %s, total iterations: %s, failing: %s" % (str(type(pool)), total, ("at index %d" % when if when else "no"))

    try:
        for res in pool.imap(prn, exception_throwing_generator(total, when), 2):
            print res
    except Exception as e:
        print "Exception from imap: %s" % str(e)

    try:
        for res in pool.imap_unordered(prn, exception_throwing_generator(total, when), 2):
            print res
    except Exception as e:
        print "Exception from imap_unordered: %s" % str(e)

def main():
    threadPool = ThreadPool(1)
    processPool = ProcessPool(1)
    test(threadPool, total=10, when=None)
    test(processPool, total=10, when=None)
    test(threadPool, total=10, when=5)
    test(processPool, total=10, when=5)
    test(threadPool, total=10, when=0)
    test(processPool, total=10, when=0)

if __name__ == "__main__":
    main()
    print "Bye bye!"