Make reusable iterable out of generator

Viewed 74

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]
1 Answers

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
Related