Memento Example: Undo in a Drawing Canvas

Each stroke saves a snapshot; undo restores the last one.

Undo in a drawing program is a natural Memento, and this example maps the template onto it. Before each change, the undo history asks the canvas for a snapshot of its strokes and its pen colour. Undo hands the last snapshot back, and the canvas restores both.

The template's parts in the canvas example
Template	Canvas example
Originator	`Canvas`, with its strokes and pen colour
Memento	`CanvasSnapshot`, a frozen dataclass
Caretaker	`UndoHistory`
save and restore	`snapshot()` and `restore(snapshot)`

from dataclasses import dataclass


@dataclass(frozen=True)
class CanvasSnapshot:
    strokes: tuple[str, ...]
    colour: str


class Canvas:
    def __init__(self) -> None:
        self._strokes: list[str] = []
        self._colour = "black"

    def draw(self, stroke: str) -> None:
        self._strokes.append(f"{self._colour} {stroke}")

    def set_colour(self, colour: str) -> None:
        self._colour = colour

    def snapshot(self) -> CanvasSnapshot:
        return CanvasSnapshot(tuple(self._strokes), self._colour)

    def restore(self, snapshot: CanvasSnapshot) -> None:
        self._strokes = list(snapshot.strokes)
        self._colour = snapshot.colour

    def show(self) -> None:
        print(f"  canvas: {self._strokes}, pen {self._colour}")


class UndoHistory:
    def __init__(self, canvas: Canvas) -> None:
        self._canvas = canvas
        self._snapshots: list[CanvasSnapshot] = []

    def before_change(self) -> None:
        self._snapshots.append(self._canvas.snapshot())

    def undo(self) -> None:
        if self._snapshots:
            self._canvas.restore(self._snapshots.pop())


canvas = Canvas()
history = UndoHistory(canvas)

history.before_change()
canvas.draw("circle")
history.before_change()
canvas.set_colour("red")
canvas.draw("line")
print("after drawing:")
canvas.show()
history.undo()
print("after one undo:")
canvas.show()
history.undo()
print("after two undos:")
canvas.show()

OutputRan with python:3.12-slim (Python 3.12.13)

after drawing:
  canvas: ['black circle', 'red line'], pen red
after one undo:
  canvas: ['black circle'], pen black
after two undos:
  canvas: [], pen black

Undo restores the pen colour as well as the strokes.

After a black circle and a red line, one undo removes the line and sets the pen back to black, because the snapshot taken before the second change held both. A second undo empties the canvas.

The snapshot copies the strokes into a tuple, so a later stroke cannot change a snapshot already taken. That copy is essential: a memento that held a reference to the canvas's own list would change with it, and every undo would restore the current state. frozen=True makes the snapshot's fields read-only too.

UndoHistory never looks at a stroke or a colour. If the canvas later saved a brush size as well, only Canvas and CanvasSnapshot would change.

← → move through sections and cards