parmap.py

import multiprocessing


def __parallel_handle(f, q_in, q_out):
    while True:
        i, x = q_in.get()
        if i is None:
            break
        q_out.put((i, f(x)))


def parmap(f, X, max_workers=multiprocessing.cpu_count(), use_tqdm=True):
    if use_tqdm:
        from tqdm import tqdm
    else:

        def tqdm(x, **params):
            return x

    if max_workers == 0:
        max_workers = multiprocessing.cpu_count()
    q_in = multiprocessing.Queue(1)
    q_out = multiprocessing.Queue()

    proc = [
        multiprocessing.Process(target=__parallel_handle, args=(f, q_in, q_out))
        for _ in range(max_workers)
    ]
    for p in proc:
        p.daemon = True
        p.start()

    try:
        sent = [
            q_in.put((i, x))
            for i, x in enumerate(tqdm(list(X), dynamic_ncols=True, smoothing=0))
        ]
        [q_in.put((None, None)) for _ in range(max_workers)]
        res = [q_out.get() for _ in range(len(sent))]
        [p.join() for p in proc]
    except KeyboardInterrupt:
        q_in.close()
        q_out.close()
        raise
    return [x for i, x in sorted(res)]