1"""
2Compare the speed of primes sequentially vs. using futures.
3"""
4
5import sys
6import time
7import math
8
9try:
10    from concurrent.futures import ThreadPoolExecutor
11except ImportError:
12    ThreadPoolExecutor = None
13try:
14    from concurrent.futures import ProcessPoolExecutor
15except ImportError:
16    ProcessPoolExecutor = None
17
18from mpi4py.futures import MPIPoolExecutor
19
20PRIMES = [
21    112272535095293,
22    112582705942171,
23    112272535095293,
24    115280095190773,
25    115797848077099,
26    117450548693743,
27    993960000099397,
28]
29
30def is_prime(n):
31    if n % 2 == 0:
32        return False
33    sqrt_n = int(math.floor(math.sqrt(n)))
34    for i in range(3, sqrt_n + 1, 2):
35        if n % i == 0:
36            return False
37    return True
38
39def sequential():
40    return list(map(is_prime, PRIMES))
41
42def with_thread_pool_executor():
43    if not ThreadPoolExecutor: return None
44    with ThreadPoolExecutor(4) as executor:
45        return list(executor.map(is_prime, PRIMES))
46
47def with_process_pool_executor():
48    if not ProcessPoolExecutor: return None
49    with ProcessPoolExecutor(4) as executor:
50        return list(executor.map(is_prime, PRIMES))
51
52def with_mpi_pool_executor():
53    with MPIPoolExecutor(4) as executor:
54        return list(executor.map(is_prime, PRIMES))
55
56def main():
57    for name, fn in [('sequential', sequential),
58                     ('threads', with_thread_pool_executor),
59                     ('processes', with_process_pool_executor),
60                     ('mpi4py', with_mpi_pool_executor)]:
61        sys.stdout.write('%s: ' % name.ljust(11))
62        sys.stdout.flush()
63        start = time.time()
64        result = fn()
65        if result is None:
66            sys.stdout.write(' not available\n')
67        elif result != [True] * len(PRIMES):
68            sys.stdout.write(' failed\n')
69        else:
70            sys.stdout.write('%5.2f seconds\n' % (time.time() - start))
71        sys.stdout.flush()
72
73if __name__ == '__main__':
74    main()
75