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(
29    func, iterable, processes, should_print_progress, filter_=None, *args, **kwargs
30):
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 the
35    results. If `processes` is greater than one, a process pool is used to run
36    the functions in parallel. `should_print_progress` is a boolean value that
37    indicates whether a string 'N of M' should be printed to indicate how many
38    of the functions have finished being run.
39    """
40    global _current
41    global _total
42    _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 result
64