Restart Generator that Consumes a Refillable Iterator

Viewed 330

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.

3 Answers
Related