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