generator for a, b, aa, ab, ba, bb, aaa, aab

Question

This:

import itertools

gen1 = itertools.product('ab', repeat=3)
for i in gen1:
    print(''.join(i), end=' ')

prints this:

aaa aab aba abb baa bab bba bbb 

but I also want it to generate a b aa ab ba bb. I'm able to do that with this:

import itertools

for i in range(1, 4):
    gen2 = itertools.product('ab', repeat=i)
    for j in gen2:
        print(''.join(j), end=' ')

which prints this:

a b aa ab ba bb aaa aab aba abb baa bab bba bbb 

Is there a more elegant way of getting this output with a single generator? Something like:

import itertools

gen3 = (your code here)
for i in gen3:
    # prints forever in the pattern: a b aa ab ba bb aaa aab aba abb baa bab bba bbb
    print(''.join(i), end=' ')

Answer 1

I would chain multiple iterators together:

from itertools import chain, product


gens = chain.from_iterable(product('ab', repeat=i) for i in range(1,4))
for i in gens:
  print(''.join(i), end=' ')

If need be, this can be extended to the infinite sequence of values by replacing range(1,4) with itertools.count(1).

If your "prints forever" comment means you want to return to a after bbb, then wrap the chained iterators in itertools.cycle.