[KongchangAI]
· 2 min read· 1,480 words

Making Python Undo Functionality Surprisingly Simple with the Memento Pattern

Making Python Undo Functionality Surprisingly Simple with the Memento Pattern

Use the Memento Pattern with Protocol and generics to implement clean undo/redo via state snapshots in Python.

The Memento Pattern implements undo by capturing object state snapshots before changes, eliminating the need to write reverse logic for every operation. Using a Python workflow as an example, the article demonstrates frozen dataclasses as immutable mementos, Protocol-based Snapshotable abstraction, and a generic History class managing the snapshot stack. It also covers trade-offs: snapshots should only capture relevant state, cannot undo external side effects, and may have memory costs — plus a comparison with the Command Pattern.

When you need to implement an "undo" feature for an object that users can modify, the most intuitive approach is to write reverse logic for every operation: undo a rename, undo adding a step, undo moving a step… Worse yet, every new operation requires a matching reverse implementation. This kind of coupling spirals out of control as features grow. The Memento Pattern offers a more elegant path — it doesn't care about the operations themselves, only about the object's state.

This article is based on a video walkthrough of Python design patterns, breaking down the Memento Pattern's implementation approach and what it teaches us about state, encapsulation, and abstraction.

From Reverse Logic to State Snapshots

Suppose we have a Workflow in an automation system that contains a name and several Step objects, each with a name and an enabled flag. Available operations include renaming the workflow, adding steps, and so on.

Implementing undo with reverse logic requires two things for every mutation: the reverse operation itself, and enough information to reconstruct the previous state. The latter is usually the most painful part. The Memento Pattern takes a fundamentally different approach: before making a change, take a snapshot of the current state. When you need to undo, simply restore the object to that snapshot.

Creating a snapshot to save the state before modifications

The first step to implementing snapshots is representing them with an immutable data class. The video uses a frozen=True dataclass named WorkflowSnapshot, containing a name field and a tuple (not a list) of steps — because the snapshot itself should never be modified. This immutable object is what the pattern calls the "Memento."

Letting Objects Own Their Save and Restore Logic

Next, we extend the Workflow class with two methods: snapshot() returns a WorkflowSnapshot capturing the current name and steps tuple; restore() accepts a snapshot and writes the data back to the object, converting the tuple back into a list.

@dataclass(frozen=True)
class WorkflowSnapshot:
    name: str
    steps: tuple[Step, ...]

class Workflow:
    def snapshot(self) -> WorkflowSnapshot:
        return WorkflowSnapshot(self.name, tuple(self.steps))

    def restore(self, snap: WorkflowSnapshot) -> None:
        self.name = snap.name
        self.steps = list(snap.steps)

With this in place, the entire undo flow is simply: take a snapshot → edit → restore the snapshot. No reverse logic needed for any operation. Even if you later add something like "shuffle step order," the snapshot and restore mechanism remains completely unaffected.

Two of the three Memento Pattern roles are now complete

There's a key design principle at work here: the workflow object owns its state, so it should also know how to snapshot and restore itself. The history management mechanism doesn't need to understand how the workflow works internally — it only needs to know that it can snapshot and restore. This is the decoupling that encapsulation provides.

Three Roles and Abstraction Expressed Through Protocol

At this point, two of the three classic Memento Pattern roles have appeared: Workflow is the Originator — the object whose state needs to be saved and restored; WorkflowSnapshot is the Memento — the carrier of the saved state. What's missing is the Caretaker — something responsible for holding onto mementos for later restoration.

These names come from the Gang of Four (GoF) design patterns book published in 1994. But as the video emphasizes, in modern Python we don't need to copy that class diagram verbatim. The Caretaker doesn't need to know Workflow exists — it just needs "something that can create and restore snapshots." This capability can be expressed with a Protocol:

from typing import Protocol

class Snapshotable[T](Protocol):
    def snapshot(self) -> T: ...
    def restore(self, snap: T) -> None: ...

Thanks to Protocol's structural subtyping, Workflow automatically satisfies this protocol without inheriting from anything. We can then write a generic History caretaker:

class History[T]:
    def __init__(self) -> None:
        self.snapshots: list[T] = []

    def save(self, obj: Snapshotable[T]) -> None:
        self.snapshots.append(obj.snapshot())

    def undo(self, obj: Snapshotable[T]) -> bool:
        if not self.snapshots:
            return False
        obj.restore(self.snapshots.pop())
        return True

undo pops the latest snapshot and returns True

History has no idea how workflows operate — it only knows "give me something that can create and restore snapshots of type T." This is exactly the right way to examine classic design patterns: don't mechanically copy them, but look at what language features are available to express the pattern more cleanly.

The Gang of Four refers to Erich Gamma, Richard Helm, Ralph Johnson, and John Vlissides, the four authors who published Design Patterns: Elements of Reusable Object-Oriented Software in 1994. The book documents 23 classic design patterns, including the Memento Pattern. The original class diagrams were designed around strongly-typed C++ and Java idioms, relying on inheritance and interface declarations to constrain role relationships. Python's Protocol uses structural subtyping (also called the static version of duck typing): a class is considered to satisfy a protocol as long as it has the required method signatures, without any explicit inheritance or registration. This contrasts with Java's implements interface declaration, which requires explicit declaration of intent. Since Python 3.8 introduced typing.Protocol, this "implicit satisfaction" can be recognized by static type checkers like mypy and pyright, achieving both flexibility and type safety.

Extending to Redo and Real-World Boundaries

Once you have a snapshot system, adding "redo" functionality is almost just a matter of adding another stack. History maintains both undo and redo lists: when performing an undo, push the snapshot onto the redo stack; when performing a redo, push it back onto the undo stack — keeping both in sync.

After undo, executing redo restores the edited state

But this pattern isn't without trade-offs:

  • What state should go into a snapshot? If a workflow also contains runtime information (like last_run_at or an execution count), should undoing a rename roll those values back too? Usually not. Snapshots shouldn't automatically serialize every attribute — they should only capture the state relevant to the restoration purpose.
  • Some state is hard to copy. Real objects may hold database connections, locks, services, caches, file handles, or large amounts of unrelated state. An explicit snapshot should clearly define the state boundary — "the state needed to restore this object for this purpose."
  • Restoring state is not the same as undoing side effects. Restoring an object cannot unsend an email, reverse an API call, or refund money to someone's account. Once external systems are involved, you need more complex architectural mechanisms like transactions and compensating operations.
  • Memory cost. Large state combined with many snapshots can create memory pressure.

When to Use Memento vs. Command Pattern

Memento is often compared to the Command Pattern, and the fundamental difference is: Memento stores state, Command stores operations.

Memento is more appropriate when state copying is cheap, you have many different types of mutations, and "restoring old state" is simpler than "reversing each operation." Command Pattern tends to be better when snapshots would be very large, individual changes are small, or the operations themselves (like "rename workflow" or "delete step") are meaningful domain concepts that need to appear in audit logs. The two can also be combined — Command records operations, Memento periodically creates checkpoints — applicable to scenarios like draft states, simulations, and partial rollback after multi-step operation failures.

One final note: if you're dealing with a very small object, just copying it directly is enough. There's no need to introduce a pile of base classes just to use the Originator, Memento, and Caretaker names from the Gang of Four. Everything shown in this article amounts to a few frozen dataclasses, one Protocol, and a small generic History class.

The truly valuable design insight from this pattern is: let the object that owns the state be responsible for capturing and restoring it, and the rest of the pattern's infrastructure only depends on that capability. In Python, with a few dataclasses, a Protocol, and modern generic syntax, you can build a remarkably clean solution.

The Command Pattern encapsulates "operations" themselves as objects, where each operation object typically includes execute() and undo() methods, and undoing calls the corresponding reverse logic. This means the system needs to maintain a queue of operation objects, and each type of operation needs its own reverse logic implemented separately. The advantage of Command is that operations become first-class citizens — they can be queued, deferred, serialized and transported across processes, and are naturally suited for audit log scenarios, because the log records "who did what and when" rather than "what the state was at the time." In a version control analogy, Memento is more like "saving a complete snapshot" while Command Pattern is more like "recording diffs" — the former has a fixed restore cost, while the latter is more storage-efficient but has a longer replay chain. When combining both patterns, Command handles fine-grained operation recording while Memento periodically creates checkpoints, enabling fast jumps to any historical node in a long operation chain.

#Python Memento Pattern#Python undo redo implementation#design patterns Python#frozen dataclass snapshot#Python Protocol structural subtyping#Command Pattern vs Memento#Python generic History class#Gang of Four Python
Share:

Related articles