Restart Generator that Consumes a Refillable Iterator

Question

I have a problem using a generator that is using a refillable iterator.

Here is my simple generator:

def hi(iterable):
  for val in iterable:
    yield val

The iterable that I pass into the hi generator is the Reservoir class from the functional_pipes repo that can be refilled after it has exhausted its elements.

I would like to consume the hi generator until StopIteration is raised and then refill the iterable and then consume it again like

refillable = Reservoir((1, 2, 3, 4))
hi_iter = hi(refillable)

print(tuple(hi_iter))

refillable((5, 6, 7, 8))
print(tuple(hi_iter))

but this prints

(1, 2, 3, 4)
()

The second tuple should also be (5, 6, 7, 8).

The only solution that I have found for this is wrapping the hi generator with a class

def super_gener(function):
  class wrapper_class:
    def __init__(self, iterable):
      self.iterable = iterable
      self.zipped = None

    def __iter__(self):
      return self

    def __next__(self):
      try:
        return next(self.zipped)

      except TypeError:
        self.zipped = function(self.iterable)
        return next(self)

      except StopIteration as err:
        self.zipped = None
        raise err

  return wrapper_class

hi_iter = super_gener(hi)(refillable)

print(tuple(hi_iter))
refillable(data)
print(tuple(hi_iter))

This solution seems a bit excessive and I'm looking for a simpler solution. Thanks for your help.

In response to Ptank: I cannot save the iterable to a tuple because the iterable is not always yielding the same items and the items are not known before refillable is filled the second time.

Answer 1

I'm afraid the only solution might be to create a refillable generator wrapper class. EDIT: Original untested code wasn't working. I have now refactored the idea below and tested it.

This object will raise StopIteration ONCE, after which it will restart. It is intended to be used with the Resettable decorator, which adds a _func attribute to the class. It should have all the same functionality of the original generator.

class ResettableGenerator():
    '''Generator wrapper that is resettable.'''
    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs
        self.start()
    def __next__(self):
        n = self.send(None)
        return n
    def __iter__(self):
        yield from self._gen
    def start(self):
        self._gen = self._func(*self.args, **self.kwargs)
    def send(self, *args, **kwargs):
        try:
            n = self._gen.send(*args, **kwargs)
            return n
        except StopIteration:
            self.start()
            raise
    def throw(self, *args, **kwargs):
        self._gen.throw(*args, **kwargs)
    def close(self):
        self._gen.close()

Here is the decorator:

def Resettable(some_func):
    cls = type(some_func.__name__, (ResettableGenerator,), {})
    cls._func = staticmethod(some_func)
    return cls

Use it like this:

@Resettable
def f():
    yield 1

Now you can do things like this:

>>> g=f()
>>> next(g)
1
>>> next(g)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 8, in __next__
  File "<stdin>", line 16, in send
StopIteration
>>> next(g)
1 #  generator has restarted itself