1import sys
2import multiprocessing
3
4
5_current = None
6_total = None
7
8
9def _init(current, total):
10    global _current
11    global _total
12    _current = current
13    _total = total
14
15
16def _wrapped_func(func_and_args):
17    func, argument, should_print_progress, filter_ = func_and_args
18
19    if should_print_progress:
20        with _current.get_lock():
21            _current.value += 1
22        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(func, iterable, processes, should_print_progress, filter_=None, *args, **kwargs):
29    """
30    A parallel map function that reports on its progress.
31
32    Applies `func` to every item of `iterable` and return a list of the
33    results. If `processes` is greater than one, a process pool is used to run
34    the functions in parallel. `should_print_progress` is a boolean value that
35    indicates whether a string 'N of M' should be printed to indicate how many
36    of the functions have finished being run.
37    """
38    global _current
39    global _total
40    _current = multiprocessing.Value('i', 0)
41    _total = multiprocessing.Value('i', len(iterable))
42
43    func_and_args = [(func, arg, should_print_progress, filter_) for arg in iterable]
44    if processes == 1:
45        result = list(map(_wrapped_func, func_and_args, *args, **kwargs))
46    else:
47        pool = multiprocessing.Pool(initializer=_init,
48                                    initargs=(_current, _total,),
49                                    processes=processes)
50        result = pool.map(_wrapped_func, func_and_args, *args, **kwargs)
51        pool.close()
52        pool.join()
53
54    if should_print_progress:
55        sys.stdout.write('\r')
56    return result
57