#!/usr/bin/env python
# coding: utf-8

# In[1]:


import ipyparallel as ipp
rc = ipp.Client()
view = rc[0]


# Assuming we have a generator function that does some work and yields results:

# In[2]:


def p():
    """My work task, yields squares up to 10"""
    import time
    for i in range(11):
        time.sleep(1)
        yield i * i


# In[3]:


for ii in p():
    print(ii)


# We can make a new generator that runs the same function on a remote engine:

# In[4]:


def remote_generator(view, f, *args, **kwargs):
    """Run a generator function remotely, returning a generator that yields in the same way."""
    # First, create the iterator by calling our function remotely
    def _create_iterator(f, *a, **kw):
        """Creates an iterator, storing it in the interactive namespace as _iterator"""
        g = globals()
        generator = f(*a, **kw)
        g['_iterator'] = iter(generator)
    
    # Create the remote iterator
    view.apply_sync(_create_iterator, f, *args, **kwargs)
    
    # then iterate over it remotely by applying the `next` builtin
    r_iterator = ipp.Reference('_iterator')
    while True:
        try:
            yield view.apply_sync(next, r_iterator)
        except ipp.RemoteError as e:
            if e.ename == 'StopIteration':
                raise StopIteration()
            else:
                raise


# In[5]:


generator = remote_generator(view, p)
generator


# In[6]:


for ii in generator:
    print(ii)