Search code examples
pythongeneratorpython-decoratorsiterable

Make reusable iterable out of generator


I want to convert generators to reusable iterables. For example, consider generator:

def myrange(n):
    for i in range(n):
        yield i

It generates one-use iterator:

x = myrange(3)
print(list(x)) # [0, 1, 2]
print(list(x)) # []

I want to add a decorator to the definition of myrange such that it produces reusable iterable (like the usual range) instead of one-use iterator:

x = myrange(3)
print(list(x)) # [0, 1, 2]
print(list(x)) # [0, 1, 2]

Solution

  • This should work:

    def mk_reusable(f):
        """
        Makes a reusable iterable out of generator by remembering its arguments
        """
    
        class MyIterable:
            def __init__(self, *args, **kwargs):
                self._args = args
                self._kwargs = kwargs
    
            def __iter__(self):
                yield from f(*self._args, **self._kwargs)
    
        return MyIterable
    
    
    # TEST
    @mk_reusable
    def myrange(n):
        for i in range(n):
            yield i
    
    x = myrange(3)
    assert list(x) == [0, 1, 2]
    assert list(x) == [0, 1, 2]
    # can consume it twice