I would like to know how to write python function which can flatten generator which yields another generators or iteables (which can also yield another generators/iterables ... possibly infinitely).
Here is example:
gen(gen(1,2,3), gen(4,5,6), [7,8,9], [gen(10,11,12), gen(13,14,15)])
note: gen
- means generator object, content between parentheses after gen
is data that will generator gen
yield.
The expected result after "flattening":
gen(1,2,3,4,5,6,7,8,9,10,11,12,13,14,15)
It is necessary for the flatten function to return generator too! (because otherwise, the preceeding usage of generators would be meaningless).
Just to note, I am using python 3.
Thanks!
The easiest way is a recursive flattening function. Assuming you want to descend into every iterable except for strings, you could do this:
def flatten(it):
for x in it:
if (isinstance(x, collections.Iterable) and
not isinstance(x, str)):
for y in flatten(x):
yield y
else:
yield x
Starting from Python 3.3, you can also write
def flatten(it):
for x in it:
if (isinstance(x, collections.Iterable) and
not isinstance(x, str)):
yield from flatten(x)
else:
yield x
The non-recursive method is essentially an unrolling of the recursive method, using a stack:
def flatten(it):
stack = []
it = iter(it)
while True:
try:
x = next(it)
except StopIteration:
if stack:
it = stack.pop()
continue
else:
return
if isinstance(x, collections.Iterable) and not isinstance(x, str):
stack.append(it)
it = iter(x)
else:
yield x