distribute

"""
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