Tower of Hanoi, ASCII animation with ANSI escape codes

Question

A few days ago I saw a nice video on YouTube with prof Thorsten Altenkirch about recursion and the Tower of Hanoi puzzle. Today I tried to reproduce the code from the video and came to this:

A = "A"
B = "B"
C = "C"

def move_one(src, dst):
    print("Move piece from {} to {}.".format(src,dst))

def move_multiple(n, src, dst, helper):
    if n == 1:
        move_one(src, dst)
    else:
        move_multiple(n-1, src, helper, dst)
        move_one(src, dst)
        move_multiple(n-1, helper, dst, src)

def main():
    move_multiple(4, A, B, C)

if __name__ == "__main__":
    main()

This works fine (although I missed the nice if n == 0 then pass) but it still is a little disappointing that we do not have a robot to actually move discs and instead use a print statement to indicate the action.

As a compromise I thought it would be cool to change the print statements to some kind of ASCII animation. I started with a print_world function and print_disc, erase_disc functions that uses ANSI escape codes to create and change the ASCII Hanoi tower world (Works on Linux, mileage on Windows may vary).

After a few hours (with trial an error) my code changed to this:

from time import sleep
import sys

A = "A"
B = "B"
C = "C"

N = 6

UP    = "\u001b[A"
DOWN  = "\u001b[B"
RIGHT = "\u001b[C"
LEFT  = "\u001b[D"

LEAD  = "  "
SPC   = " "

# initial state
state = {
          A : list(range(N,0,-1)),
          B : [],
          C : [],
        }

def change_state(src,dst):
    global state
    state[dst].append(state[src].pop())

def print_world():
    """Print empty Hanoi world, the poles and labels"""
    print("\n\n")
    print((LEAD + SPC.join([SPC * N + "|" + SPC * N] * 3) + "\n") * N)
    print(LEAD + SPC * N + (SPC * (2 * N + 1)).join([A,B,C]))

def print_disc(n, up, right):
    """Print disc of certain size at certain screen coordinates"""
    tkn = str(n)
    print (   UP * up 
            + RIGHT * right
            + SPC * (N-n) + tkn * n + RIGHT + tkn * n + SPC * (N-n)
            + LEFT * (2 * N + 1)
            + LEFT * right
            + DOWN * up,
           end = "")

def erase_disc(up, right):
    """Erase disc at certain screen coordinates"""
    print (   UP * up 
            + RIGHT * right
            + SPC * N + RIGHT + SPC * N
            + LEFT * (2 * N + 1)
            + LEFT * right
            + DOWN * up,
           end = "")

def find_screen_pos(pole, height):
    up = 3 + height
    height = len(LEAD) + (2 * N + 2) * {A:0, B:1, C:2}[pole]
    return up, height

def print_state():
    global state
    for pole, discs in state.items():
        for height, disc in enumerate(discs):
            up, right = find_screen_pos(pole, height)
            print_disc(disc, up, right)
    sys.stdout.flush()

def find_screen_path(pole_1, height_1, pole_2, height_2):
    cur_up, cur_right = find_screen_pos(pole_1, height_1)
    end_up, end_right = find_screen_pos(pole_2, height_2)
    direction = 0
    if end_right > cur_right : direction = 1
    if end_right < cur_right : direction = -1
    path = [(cur_up, cur_right)]
    while cur_up < N + 4:
        cur_up += 1
        path.append((cur_up, cur_right))
    while cur_right != end_right:
        cur_right += direction
        path.append((cur_up, cur_right))
    while cur_up != end_up:
        cur_up -= 1
        path.append((cur_up, cur_right))
    return path

def animate_move(src, dst):
    height_src = len(state[src])-1
    height_dst = len(state[dst])
    n = state[src][-1]
    path = find_screen_path(src, height_src, dst, height_dst)
    for old, new in zip(path[:-1], path[1:]):
        erase_disc(*old)
        print_disc(n, *new)
        sys.stdout.flush()
        sleep(0.01)

def move_one(src, dst):
    animate_move(src, dst)
    change_state(src, dst)
    sleep(0.2)

def move_multiple(n, src, dst, helper):
    if n == 1:
        move_one(src, dst)
    else:
        move_multiple(n-1, src, helper, dst)
        move_one(src, dst)
        move_multiple(n-1, helper, dst, src)

def main():
    print_world()
    print_state()
    sleep(2.0)
    while(True):
        move_multiple(N,A,B,C)
        sleep(2.0)
        move_multiple(N,B,C,A)
        sleep(2.0)
        move_multiple(N,C,A,B)
        sleep(2.0)


if __name__ == "__main__":
    main()

This relatively short piece of code (126 lines) works and at this moment I understand it, but do you, and will I in six month time? What should I change to the code to make it more readable and understandable (I know I've sinned against PEP8 sometimes).

Answer 1

docstrings and comments

There are lots of formulas in the drawing code, like:

+ SPC * (N-n) + tkn * n + RIGHT + tkn * n + SPC * (N-n)

How did you arrive at those. How would they change if you wanted to do some other kind of animation? I would add a ASCII art drawing of a puzzle being solved, annotated with the dimensions and formulas.

Hidden interfaces

Hidden interfaces are bad. move_multiple() takes a parameter n that is the number of disks to move. This implies that you can select what size problem to solve by changing n. However, the drawing/animation functions are coded to use a global variable N for the number of discs. And N is hard coded to 6 disks.

separate concerns

The function move_one() does three things: 1) starts the animation for a move, 2) helps solve the problem (updating the state), and 3) adds a delay between steps in the animation. That can make it harder to maintain or reuse the code. For example, a year from now, you want to change the delay. Will you remember the delay is in the tower solving code and not somewhere in the animation code? Or you want to modify the code to drive an HTML canvas routine. Will you remember that each animation step is started from move_one()?

It would be better if the problem solving code just solved the problem and the animation code just did the animation. Modify the problem solving code to return a list of moves, or, better yet, turn it into a generator that yields the steps as needed:

INTER_STEP_PAUSE = 0.2
INTER_PUZZLE_PAUSE = 2.0

# labels for the towers.  TODO: internationalize
TOWER1 = 'A'
TOWER2 = 'B'
TOWER3 = 'C'

def hanoi(n, src, dst, helper):
    """A generator to solve an n-disk Tower of Hanoi problem.

    It yields the steps for moving a stack of n discs from the source tower
    to the target tower. `helper` is the other tower.

    Yields (from, to) tuples, which means move the top disc from `from` to `to`.

    Can be used like:

        for src, dst in hanoi(3, 'Alpha', 'Bravo', 'Charlie'):
            print(f"Move disk from {src} to {dst}")
    """

    if n == 1:
        yield (source, target)

    else:
        hanoi(n-1, source, helper, target)
        hanoi( 1,  source, target, helper)
        hanoi(n-1, helper, target, source)

Then main() can be something like:

def main():
    n = int(sys.argv[1])

    print_world(n)
    print_state(n)
    sleep(2.0)

    tower1, tower2, tower3 = LABEL1, LABEL2, LABEL3

    while True:
        for source,target in hanoi(n, tower1, tower2, tower3):
            animate(source, target)
            sleep(INTER_STEP_PAUSE)

        # rotate the tower and solve it again
        tower1, tower2, tower3 = tower2, tower3, tower1
        sleep(INTER_PUZZLE_PAUSE)

Implementing the display and animation code is left as an exercise for the reader.