Python: translating a printing recursive function into a generator

Question

I have found this function:

def hanoi(pegs, start, target, n):
    assert len(pegs[start]) >= n, 'not enough disks on peg'
    if n == 1:
        pegs[target].append(pegs[start].pop())
        print '%i -> %i: %s' % (start, target, pegs)
    else:
        aux = 3 - start - target  # start + target + aux = 3
        hanoi(pegs, start, aux, n-1)
        hanoi(pegs, start, target, 1)
        hanoi(pegs, aux, target, n-1)

At Trying to implement recursive Tower of Hanoi algorithm with arrays here on StackOverflow.

Now I need to modify this so that instead of the start, target, pegs being printed, the pegs variable is yielded out each iteration.

For example I expect this output out of the new function (pretty printed):

>>> list( hanoi([[120, 90, 60, 30], [], []]) )
[ [[120, 90, 60], [30], []],
  [[120, 90], [30], [60]],
  [[120, 90], [], [60, 30]],
  [[120], [90], [60, 30]],
  [[120, 30], [90], [60]],
  [[120, 30], [90, 60], []],
  [[120], [90, 60, 30], []],
  [[], [90, 60, 30], [120]],
  [[], [90, 60], [120, 30]],
  [[60], [90], [120, 30]],
  [[60, 30], [90], [120]],
  [[60, 30], [], [120, 90]],
  [[60], [30], [120, 90]],
  [[], [30], [120, 90, 60]],
  [[], [], [120, 90, 60, 30]],
]

This is how I tried to modify it:

def hanoi(pegs, start, target, n):
    assert len(pegs[start]) >= n, 'not enough disks on peg'
    if n == 1:
        pegs[target].append(pegs[start].pop())
        yield pegs
    else:
        aux = 3 - start - target  # start + target + aux = 3
        hanoi(pegs, start, aux, n-1)
        hanoi(pegs, start, target, 1)
        hanoi(pegs, aux, target, n-1)

But, (input pegs numbers are a bit bigger because of graphical purposes):

>>> pegs = [[120, 90, 60, 30], [], []]
>>> print(list(hanoi(pegs, 0, 2, 4)))
[]

The output is just an empty list.

Trying to make copies of the lists via [:] did not help, and I am left very confused, maybe print can always print out, but yield gets "stuck" inside deep recursion levels, so it is yielded to less deep recursion and not to the outside. Also using a list with append does not work:

def hanoi(pegs, start, target, n):
    assert len(pegs[start]) >= n, 'not enough disks on peg'
    out = []
    if n == 1:
        pegs = pegs[:]
        pegs[target].append(pegs[start].pop())
        out.append( pegs )
    else:
        aux = 3 - start - target  # start + target + aux = 3
        hanoi(pegs, start, aux, n-1)
        hanoi(pegs, start, target, 1)
        hanoi(pegs, aux, target, n-1)
    return out

I also tried following advice from Python: using a recursive algorithm as a generator

def hanoi(pegs, start, target, n):
    assert len(pegs[start]) >= n, 'not enough disks on peg'
    if n == 1:
        pegs = pegs[:]
        pegs[target].append(pegs[start].pop())
        yield pegs
    else:
        aux = 3 - start - target  # start + target + aux = 3
        for i in hanoi(pegs, start, aux, n-1): yield i
        for i in hanoi(pegs, start, target, 1): yield i
        for i in hanoi(pegs, aux, target, n-1): yield i

By yielding from nested for loops, but it failed.

How can I write such generator (that I need for graphical purposes)?

The generator will be used like this:

pegs = [[120, 90, 60, 30], [], []]
positions = hanoi(pegs, 0, 2, 4)

for position in positions:
    screen.fill((255, 255, 255)) 

    print(index, peg_history[index])
    for i, pegs in enumerate(position):
        display_pegs(pegs, 100 + 180*i, 300, screen)
    pygame.display.update()
    time.sleep(0.5)

Answer 1

a generator version could look like this:

def hanoi_yield(pegs, start, target, n):
    # pegs will be modified!
    assert len(pegs[start]) >= n, 'not enough disks on peg'

    if n == 1:
        pegs[target].append(pegs[start].pop())
        yield pegs
    else:
        aux = 3 - start - target  # start + target + aux = 3
        yield from hanoi_yield(pegs, start, aux, n-1)
        yield from hanoi_yield(pegs, start, target, 1)
        yield from hanoi_yield(pegs, aux, target, n-1)

pegs = [[120, 90, 60, 30], [], []]
for item in hanoi_yield(pegs, 0, 2, 4):
    print(item)

with the output:

[[120, 90, 60], [30], []]
[[120, 90], [30], [60]]
[[120, 90], [], [60, 30]]
[[120], [90], [60, 30]]
[[120, 30], [90], [60]]
[[120, 30], [90, 60], []]
[[120], [90, 60, 30], []]
[[], [90, 60, 30], [120]]
[[], [90, 60], [120, 30]]
[[60], [90], [120, 30]]
[[60, 30], [90], [120]]
[[60, 30], [], [120, 90]]
[[60], [30], [120, 90]]
[[], [30], [120, 90, 60]]
[[], [], [120, 90, 60, 30]]

the only 'trick' here is to yield from hanoi_yield because hanoi_yield is a generator.

the down-side: this returns a reference to the same list all the time and changes the input list pegs (which is just the return value)! that may not be wished or useful... more below:

---

a version that does not change the first argument (pegs) and returns an individual list every time (and can therefore be used in the list constructor). i had to add an auxiliary variable _work_pegs as the algorithm needs to change the this list. pegs is now unchanged. i also yield a deepcopy of the result (we are handling lists of lists here; a regular copy would not work):

from copy import deepcopy

def hanoi_yield(pegs, start, target, n, _work_pegs=None):

    if _work_pegs is None:
        _work_pegs = deepcopy(pegs)
        # or (this way pegs could be a tuple of tuples):
        # _work_pegs = [list(item) for item in pegs]

    assert len(_work_pegs[start]) >= n, 'not enough disksdef on peg'

    if n == 1:
        _work_pegs[target].append(_work_pegs[start].pop())
        yield deepcopy(_work_pegs)
        # or (returning tuples might be nice...):
        # yield tuple(tuple(item) for item in _work_pegs)
    else:
        aux = 3 - start - target  # start + target + aux = 3
        yield from hanoi_yield(pegs, start, aux, n-1, _work_pegs)
        yield from hanoi_yield(pegs, start, target, 1, _work_pegs)
        yield from hanoi_yield(pegs, aux, target, n-1, _work_pegs)

finally this works:

pegs = [[120, 90, 60, 30], [], []]
lst = list(hanoi_yield(pegs, 0, 2, 4))
print(lst)