Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
"""
Syntactic sugar to parallelize loops.
:author: Michael Bironneau <michael.bironneau@openenergi.com>
:date: 2014-10-29
Example::
@distribute_over(range(10))
def get_squares(i):
time.sleep(1) #Do something time-consuming
return i**2
print(get_squares())
If you need to set the `shard` parameter dynamically (that is, after function definition), use `@distribute` instead::
@distribute
def get_squares(i):
time.sleep(1) #Do something time-consuming
return i**2
print(get_squares(range(10))
"""
from concurrent.futures import ThreadPoolExecutor
def distribute_over(shards, threads=5):
"""Parallel for loop over an iterable 'shards'"""
def wrapper(fn):
def func_wrapper():
result = []
with ThreadPoolExecutor(max_workers=threads) as executor:
gen = executor.map(fn, shards)
for res in gen:
result.append(res)
return result
return func_wrapper
return wrapper
def distribute(fn, threads=5):
"""Parallel for loop over an iterable 'shards', which is passed as the first argument of fn"""
def func_wrapper(shards):
result = []
with ThreadPoolExecutor(max_workers=threads) as executor:
gen = executor.map(fn, shards)
for res in gen:
result.append(res)
return result
return func_wrapper
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment