Decotator to memoize results of generator. This will scan until end of the generator at the first of iterating.

memoize_generator.py

from functools import wraps


class MemGen:
    def __init__(self, generator):
        self.generator = generator
        self.mem = []
        self.__doc__ = generator.__doc__

    def __iter__(self):
        if not self.mem:
            self.mem = list(self.generator)
        return iter(self.mem)

    def __next__(self):
        if not self.mem:
            self.mem = list(self.generator)
        for m in self.mem:
            return m


def memoize_generator(func):
    @wraps(func)
    def _wrapped(*args, **kwargs):
        return MemGen(func(*args, **kwargs))
    return _wrapped