File size: 1,382 Bytes
26225c5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import multiprocessing
from itertools import repeat
__all__ = ['starmap_with_kwargs']
def starmap_with_kwargs(fn, args_iter, kwargs_iter, processes=4):
"""By default, starmap only accepts args and not kwargs. This is a
helper to get around this problem.
:param fn: callable
The function to starmap
:param args_iter: iterable
Iterable of the args
:param kwargs_iter: iterable or dict
Kwargs for `fn`. If an iterable is passed, the corresponding
kwargs will be passed to each process. If a dictionary is
passed, these same kwargs will be repeated and passed to all
processes. NB: this behavior only works for kwargs, if the same
args need to be passed to the `fn`, the adequate iterable must
be passed as input
:param processes: int
Number of processes
:return:
"""
# Prepare kwargs
if kwargs_iter is None:
kwargs_iter = repeat({})
if isinstance(kwargs_iter, dict):
kwargs_iter = repeat(kwargs_iter)
# Apply fn in multiple processes
with multiprocessing.get_context("spawn").Pool(processes=processes) as pool:
args_for_starmap = zip(repeat(fn), args_iter, kwargs_iter)
out = pool.starmap(apply_args_and_kwargs, args_for_starmap)
return out
def apply_args_and_kwargs(fn, args, kwargs):
return fn(*args, **kwargs)
|