zippin' zippers in Python

ziptree.py

# Language path: Haskell->Ocaml->Python, mypy typesystem [semantic loss ~31%]
# Subject: Zippers as the key insight

# what great fun! since mypy ~0.991 (?) we can finally, finally work with
#   recursive types. this opens up Python to some traditional functional
#   programming practices with a minimum of bullshit.

from typing import NamedTuple, TypeAlias

# let's walk amongst the trees for a bit.
# we can represent binary trees as nested tuples
# a tree is either:
#   - the empty tree (represented by None), or
#   - a branch with data element x and left and right children (x, left, right)

class Branch(NamedTuple):
    value: int
    left: 'Tree'
    right: 'Tree'

Tree: TypeAlias = Branch | None

# for example:

example_tree: Tree = Branch(3,
  Branch(2, Branch(1, None, None), None),
  Branch(5, Branch(4, None, None), Branch(6, None, None))
)

# we can make this easier to understand with some ASCII art:

def pprint(tree: Tree, indent: int = 0) -> None:
    if tree is None:
        print(indent * ' ' + '·')
    else:
        print(indent * ' ' + str(tree.value))
        pprint(tree.left, indent = indent + 4)
        pprint(tree.right, indent = indent + 4)

# suppose that for some contrived reason (maybe we're solving a programming
#   puzzle [https://adventofcode.com/2022/day/7]), we have a sequence of
#   commands we want to run against this tree, where each command may either
#   navigate the tree ("go left", "go right", "go back to the parent") or
#   edit the current position ("set value to 3", "replace left child with None")
# in a purely functional program, all data is immutable: we will not be editing
#   the tree in place; rather, each command must produce a new copy of the data
#   with the necessary edits applied.
# one more thing: for similarly contrived reasons, we care about the final state
#   of the ENTIRE tree (in other words, the "root") after all commands are
#   applied.

# this makes life hard! let's try to write a program for "go left, then go left,
#   then set value to 77."

def left_left_77_first_try(tree: Tree) -> Tree:
    # navigation is pretty straight-forward:
    if tree is None:
        raise ValueError("can't go left!")
    tree = tree.left
    if tree is None:
        raise ValueError("can't go left!")
    tree = tree.left
    if tree is None:
        raise ValueError("can't go left!")

    # hmm - this doesn't feel quite right!
    return tree._replace(value = 77)

# that's not right at all - we've edited the left, left subtree, but we've
#   "forgotten how we got there"; remember, we're interested in making
#   edits at the leaves but retaining the whole tree (with those edits)
# to make that work, we need to take apart the tree on our way "down" to the
#   destination, and rebuild it with edits on our way back "up".

def left_left_77_second_try(tree: Tree) -> Tree:
    if tree is None:
        raise ValueError("can't go left!")
    first_left = tree.left
    if first_left is None:
        raise ValueError("can't go left!")
    second_left = first_left.left
    if second_left is None:
        raise ValueError("can't go left!")

    # rebuild the tree, making replacements...
    return tree._replace(
      left = first_left._replace(
        left = second_left._replace(value = 77)))

# that worked! great, that was just a one-off transformation, right?
# <meme> ...right?

# oh dear. it seems the actual data is a sequence of these relative commands
commands = [
    'left', 'left', 'set 77',
    'up', 'set 99',
    'up', 'right', 'right', 'set 33'
]

# well, we COULD translate each sequence of moves followed by a set to the
#   appropriate traversal and edit, like `left_left_77_second_try` - but we'll
#   be walking and rebuilding the WHOLE tree, from the root, every time we make
#   an edit.
# that clearly won't scale to the thousands of nodes that this toy problem
#   obviously will one day reach.

# let's slow down and think a little. we've already noticed that when we walk
#   a tree to make edits, we have to remember "the rest of the tree" in order
#   to reconstruct it; we also have to know "how we got there" to know which
#   parts of the original tree to replace as we work our way back up.
# the mechanism is kind of like a zipper opening and closing: we "unzip" the
#   tree apart as we move "down" to the focus of the edit, and "zip" it back
#   together when we're done.
# well, let's take that and make it concrete!

# a zipper for a binary tree represents a focused node in the tree, with the
#   additional information required to navigate around and rebuild the whole
#   things.
# in practice, we'll model it as

from enum import Enum, auto

# the directions we can "turn" along a path

class Direction(Enum):
    Left = auto()
    Right = auto()

# a step along the path: a direction, the value at the branch,
#   and the tree which was the "road not taken"!

class PathStep(NamedTuple):
    dir: Direction
    value: int
    other: Tree

# the zipper just combines a stack of the steps we took to get to the
#   focused subtree, and the focused subtree itself!
# I'm going to use a list as a purely-functional stack by copying;
#   this is a little ugly
# we don't need "deepcopy" because all the data is immutable!

class Zipper(NamedTuple):
    path: list[PathStep]
    focus: Tree

# building a zipper for a given root tree is easy!

def to_zipper(tree: Tree) -> Zipper:
    return Zipper([], tree)

# recovering the focused tree from a zipper is also easy
def from_zipper(zipper: Zipper) -> Tree:
    return zipper.focus

# we can now build operations on zippers, for navigating the tree!

# going left or right requires remembering the value at the branch,
#   and the other subtree, pushing these onto the path
# these operations can fail if the current focus is a dead-end (None)!

def go_left(zipper: Zipper) -> Zipper | None:
    if zipper.focus is None:
        return None
    step = PathStep(Direction.Left, zipper.focus.value, zipper.focus.right)
    return Zipper([step] + zipper.path, zipper.focus.left)

def go_right(zipper: Zipper) -> Zipper | None:
    if zipper.focus is None:
        return None
    step = PathStep(Direction.Right, zipper.focus.value, zipper.focus.left)
    return Zipper([step] + zipper.path, zipper.focus.right)

# going "up" requires rebuilding the parent branch from the remembered
#   value and "other" fork; if we're at the root already, we'll just
#   return the root zipper.

def go_up(zipper: Zipper) -> Zipper:
    match zipper.path:
        case []:
            return zipper
        case [PathStep(Direction.Left, value, right), *rest]:
            return Zipper(rest, Branch(value, zipper.focus, right))
        case [PathStep(Direction.Right, value, left), *rest]:
            return Zipper(rest, Branch(value, left, zipper.focus))
    raise RuntimeError("this can't happen but mypy doesn't know it!")

# going all the way back to the root is just a matter of recursing "up"
#   until we're out of path

def go_root(zipper: Zipper) -> Zipper:
    if not zipper.path:
        return zipper
    return go_root(go_up(zipper))

# for convenience, we'll provide a higher-order function for "editing" the focus

from typing import Callable

def update_focus(zipper: Zipper, f: Callable[[Tree], Tree]) -> Zipper:
    return zipper._replace(focus = f(zipper.focus))

# let's navigate around our original tree a little, and try making an edit!

# remember our "command sequence"?
def apply_commands(commands: list[str], tree: Tree) -> Tree:
    def apply_command(z: Zipper, c: str) -> Zipper:
        match c.split():
            # we'll be "permissive" of bad commands, and just stay where we
            #   are if navigation fails
            case ['left']:
                return go_left(z) or z
            case ['right']:
                return go_right(z) or z
            case ['up']:
                return go_up(z)
            case ['set', n]:
                return update_focus(z,
                  lambda t: t and t._replace(value = int(n)))
            case _:
                raise ValueError(f'bad command: {c}')

    from functools import reduce
    return from_zipper(go_root(
      reduce(apply_command, commands, to_zipper(tree))))

def main() -> None:
    print(example_tree)

    pprint(example_tree)

    pprint(left_left_77_first_try(example_tree))

    pprint(left_left_77_second_try(example_tree))

    left = go_left(to_zipper(example_tree))
    print(left)
    assert left is not None

    left2 = left and go_left(left)
    print(left2)
    assert left2 is not None

    print(go_root(left2))

    # dumb Python trick: if `x: None | T` and `f: Callable[[T], T]`
    #   then `x and f(x): None | T`
    root_zipper = go_root(
      update_focus(left2, lambda t: t and t._replace(value = 77)))
    new_tree = from_zipper(root_zipper)
    pprint(new_tree)

    pprint(apply_commands(commands, example_tree))

if __name__ == '__main__':
    main()