Source code for monty.multiprocessing

# coding: utf-8
"""
Multiprocessing utilities.
"""

from multiprocessing import Pool
from typing import Iterable, Callable

try:
    from tqdm.autonotebook import tqdm
except ImportError:
    tqdm = None


[docs]def imap_tqdm(nprocs: int, func: Callable, iterable: Iterable, *args, **kwargs): """ A wrapper around Pool.imap. Creates a Pool with nprocs and then runs a f unction over an iterable with progress bar. :param nprocs: Number of processes :param func: Callable :param iterable: Iterable of arguments. :param args: Passthrough to Pool.imap :param kwargs: Passthrough to Pool.imap :return: Results of Pool.imap. """ if tqdm is None: raise ImportError("tqdm must be installed for this function.") data = [] with Pool(nprocs) as pool: try: n = len(iterable) # type: ignore except TypeError: n = None # type: ignore with tqdm(total=n) as pbar: for i, d in enumerate(pool.imap(func, iterable, *args, **kwargs)): pbar.update() data.append(d) return data