55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
import sys
|
|
import multiprocessing
|
|
|
|
|
|
_current = None
|
|
_total = None
|
|
|
|
|
|
def _init(current, total):
|
|
global _current
|
|
global _total
|
|
_current = current
|
|
_total = total
|
|
|
|
|
|
def _wrapped_func(func_and_args):
|
|
func, argument, should_print_progress = func_and_args
|
|
|
|
if should_print_progress:
|
|
with _current.get_lock():
|
|
_current.value += 1
|
|
sys.stdout.write('\r\t{} of {}'.format(_current.value, _total.value))
|
|
sys.stdout.flush()
|
|
|
|
return func(argument)
|
|
|
|
|
|
def pmap(func, iterable, processes, should_print_progress, *args, **kwargs):
|
|
"""
|
|
A parallel map function that reports on its progress.
|
|
|
|
Applies `func` to every item of `iterable` and return a list of the
|
|
results. If `processes` is greater than one, a process pool is used to run
|
|
the functions in parallel. `should_print_progress` is a boolean value that
|
|
indicates whether a string 'N of M' should be printed to indicate how many
|
|
of the functions have finished being run.
|
|
"""
|
|
global _current
|
|
global _total
|
|
_current = multiprocessing.Value('i', 0)
|
|
_total = multiprocessing.Value('i', len(iterable))
|
|
|
|
func_and_args = [(func, arg, should_print_progress,) for arg in iterable]
|
|
if processes <= 1:
|
|
result = map(_wrapped_func, func_and_args, *args, **kwargs)
|
|
else:
|
|
pool = multiprocessing.Pool(initializer=_init,
|
|
initargs=(_current, _total,),
|
|
processes=processes)
|
|
result = pool.map(_wrapped_func, func_and_args, *args, **kwargs)
|
|
|
|
if should_print_progress:
|
|
sys.stdout.write('\r')
|
|
return result
|