64 lines · python
1import sys2import multiprocessing3 4 5_current = None6_total = None7 8 9def _init(current, total):10 global _current11 global _total12 _current = current13 _total = total14 15 16def _wrapped_func(func_and_args):17 func, argument, should_print_progress, filter_ = func_and_args18 19 if should_print_progress:20 with _current.get_lock():21 _current.value += 122 sys.stdout.write("\r\t{} of {}".format(_current.value, _total.value))23 sys.stdout.flush()24 25 return func(argument, filter_)26 27 28def pmap(29 func, iterable, processes, should_print_progress, filter_=None, *args, **kwargs30):31 """32 A parallel map function that reports on its progress.33 34 Applies `func` to every item of `iterable` and return a list of the35 results. If `processes` is greater than one, a process pool is used to run36 the functions in parallel. `should_print_progress` is a boolean value that37 indicates whether a string 'N of M' should be printed to indicate how many38 of the functions have finished being run.39 """40 global _current41 global _total42 _current = multiprocessing.Value("i", 0)43 _total = multiprocessing.Value("i", len(iterable))44 45 func_and_args = [(func, arg, should_print_progress, filter_) for arg in iterable]46 if processes == 1:47 result = list(map(_wrapped_func, func_and_args, *args, **kwargs))48 else:49 pool = multiprocessing.Pool(50 initializer=_init,51 initargs=(52 _current,53 _total,54 ),55 processes=processes,56 )57 result = pool.map(_wrapped_func, func_and_args, *args, **kwargs)58 pool.close()59 pool.join()60 61 if should_print_progress:62 sys.stdout.write("\r")63 return result64