parallel.py 文件源码

def pmap(func, iterable, args = tuple(), kwargs = dict(), processes = 1, ntotal = None):
    """
    Parallel application of a function with keyword arguments.

    Parameters
    ----------
    func : callable
        Function to be applied to every element of `iterable`.
    iterable : iterable
        Iterable of items to be mapped.
    args : tuple, optional
        Positional arguments of `function`.
    kwargs : dictionary, optional
        Keyword arguments of `function`.
    processes : int or None, optional
        Number of processes to use. If `None`, maximal number of processes
        is used. Default is one.
    ntotal : int or None, optional
        If the length of `iterable` is known, but passing `iterable` as a list
        would take too much memory, the total length `ntotal` can be specified. This
        allows for `pmap` to chunk better.

    Yields
    ------
    Mapped values.

    See Also
    --------
    pmap_unordered : parallel map that does not preserve order

    Notes
    -----
    If `processes` is 1, `pmap` reduces to `map`, with the added benefit of
    of using `kwargs`
    """
    func = partial(func, *args, **kwargs)

    if processes == 1:
        yield from map(func, iterable)
        return

    with Pool(processes) as pool:
        chunksize = 1
        if isinstance(iterable, Sized):
            chunksize = max(1, int(len(iterable)/pool._processes))
        elif ntotal is not None:
            chunksize = max(1, int(ntotal/pool._processes))

        yield from pool.imap(func = func, iterable = iterable, chunksize = chunksize)