parmap with monitoring

parmap.py

def _parmap(fun, q_in, q_out, init=None):
    init() if init else None
    while True:
        i, x = q_in.get()
        ret = fun(x)
        q_out.put((i, ret))


def parmap(fun, tasks, n=None, init=None):
    from multiprocessing import Pool, Queue
    from tqdm import tqdm

    q_in, q_out = Queue(), Queue()
    for item in enumerate(tasks):
        q_in.put(item)
    
    with Pool(n, _parmap, (fun, q_in, q_out, init)):
        n = len(tasks)
        ret = [None] * n
        for _ in tqdm(range(n)):
            i, x = q_out.get()
            ret[i] = x
    
    return ret