57 lines
1.7 KiB
Python
57 lines
1.7 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, filter_ = 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, filter_)
|
||
|
|
|
||
|
|
|
||
|
|
def pmap(func, iterable, processes, should_print_progress, filter_=None, *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, filter_) for arg in iterable]
|
||
|
|
if processes == 1:
|
||
|
|
result = list(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)
|
||
|
|
pool.close()
|
||
|
|
pool.join()
|
||
|
|
|
||
|
|
if should_print_progress:
|
||
|
|
sys.stdout.write('\r')
|
||
|
|
return result
|