From Chaos to Clarity: A Visual Guide to the State Design Pattern in Python

Have you ever found yourself trapped in a labyrinth of if/elif/else statements, where each branch checks the status of an object before deciding what to do next? This is a common scenario when building systems with objects that have complex life cycles, like a document, an order, or a user account. As new states and transitions are added, this logic becomes brittle, hard to read, and a nightmare to maintain.

This is where the State Design Pattern comes to the rescue. It allows an object to alter its behavior when its internal state changes. The object will appear to change its class.

In this guide, we’ll refactor a messy document management system into a clean, robust, and scalable solution using the State Pattern, complete with visual diagrams to make every step clear.

The Initial Problem: A Monolithic Document Class

Let’s start with our “before” picture. We have a Document class that manages its own state with a simple string attribute. Every action requires a check on this attribute, leading to a tangled mess of conditionals.

# The "Before" - A single class handling all state logic
class Document:
    def __init__(self):
        self._state = "draft"
        print(f"Document created. Current state: '{self._state}'")

    def edit(self):
        if self._state == "draft" or self._state == "needs_revision":
            print("Editing the document...")
        else:
            print(f"Error: Cannot edit in '{self._state}' state.")

    def submit(self):
        if self._state == "draft" or self._state == "needs_revision":
            self._state = "under_review"
            print(f"Document submitted. New state: '{self._state}'")
        else:
            print(f"Error: Cannot submit in '{self._state}' state.")

    def publish(self):
        if self._state == "under_review":
            self._state = "published"
            print(f"Document published. New state: '{self._state}'")
        else:
            print(f"Error: Cannot publish from '{self._state}' state.")

    def reject(self):
        if self._state == "under_review":
            # In a real app, you might have more complex logic here
            self._state = "needs_revision"
            print(f"Document needs revision. New state: '{self._state}'")
        else:
            print(f"Error: Cannot reject from '{self._state}' state.")

While this works for a simple case, imagine adding new states like archived, expired, or final_rejection. Each new state requires modifying every single method, violating the Open/Closed Principle and increasing the risk of bugs.

The Solution: Introducing the State Design Pattern

The State Pattern proposes a cleaner structure by encapsulating state-specific logic into separate objects.

Here are the core components:

1. Context: The main object whose behavior we want to control (Document). It holds a reference to its current state object.
2. State Interface: An abstract class that defines the methods all concrete states must implement (e.g., edit, submit, publish).
3. Concrete States: Individual classes for each state (DraftState, UnderReviewState, etc.), implementing the behavior and transition logic for that specific state.

This diagram illustrates the relationship between the components:

classDiagram
    class Context {
        -State state
        +request()
        +transition_to(State)
    }
    class State {
        <<abstract>>
        #Context context
        +handle()
    }
    class ConcreteStateA {
        +handle()
    }
    class ConcreteStateB {
        +handle()
    }

    Context o-- State
    State <|-- ConcreteStateA
    State <|-- ConcreteStateB

Step 1: Define the Project Structure

First, let’s organize our files. We’ll separate the context, states, and the main execution logic.

document_project/
├── document.py      # The Context class
├── states.py        # The State interface and all Concrete States
└── main.py          # Our application's entry point

Step 2: Create the State Abstract Base Class

In states.py, we’ll define our State interface using Python’s abc module. This ensures that every state class we create will have the same set of methods.

# states.py
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING

# Forward declaration to prevent circular imports
if TYPE_CHECKING:
    from document import Document

class State(ABC):
    """
    The base State class declares methods that all Concrete State should
    implement and also provides a backreference to the Context object,
    associated with the State.
    """
    _context: Document

    def set_context(self, context: Document):
        self._context = context

    @property
    def name(self) -> str:
        return self.__class__.__name__

    @abstractmethod
    def edit(self):
        pass

    @abstractmethod
    def submit(self):
        pass

    @abstractmethod
    def publish(self):
        pass

    @abstractmethod
    def reject(self):
        pass

[!TIP] We use if TYPE_CHECKING: and a string forward reference ('Document') to give our type checker context without creating a circular import between states.py and document.py.

Step 3: Implement the Concrete States

Now for the fun part. We’ll create a class for each state. Each class contains the logic that used to be in our if/elif blocks.

DraftState

A document in the draft state can be edited or submitted.

# states.py (continued)

class DraftState(State):
    def edit(self):
        print("Editing the document while in Draft.")

    def submit(self):
        print("Submitting the document for review.")
        self._context.transition_to(UnderReviewState())

    def publish(self):
        print("Error: Cannot publish directly from Draft.")

    def reject(self):
        print("Error: Cannot reject a document that is in Draft.")

UnderReviewState

When under review, a document can’t be edited. It can only be published or rejected.

# states.py (continued)

class UnderReviewState(State):
    def edit(self):
        print("Error: Cannot edit while under review.")

    def submit(self):
        print("Error: Document is already under review.")

    def publish(self):
        print("Publishing the document.")
        self._context.transition_to(PublishedState())

    def reject(self):
        print("Rejecting the document. It needs revision.")
        self._context.transition_to(NeedsRevisionState())

Step 4: Refactor the Document (Context) Class

With the state logic extracted, our Document class becomes incredibly simple. It delegates all state-specific actions to its current state object.

Here is the “diff” view of the transformation:

--- a/document_messy.py
+++ b/document_clean.py
@@ -1,29 +1,27 @@
-# The "Before" - A single class handling all state logic
-class Document:
-    def __init__(self):
-        self._state = "draft"
-        print(f"Document created. Current state: '{self._state}'")
+from __future__ import annotations
+from states import State, DraftState
 
-    def edit(self):
-        if self._state == "draft" or self._state == "needs_revision":
-            print("Editing the document...")
-        else:
-            print(f"Error: Cannot edit in '{self._state}' state.")
+class Document:
+    _state: State = None
 
-    def submit(self):
-        if self._state == "draft" or self._state == "needs_revision":
-            self._state = "under_review"
-            print(f"Document submitted. New state: '{self._state}'")
-        else:
-            print(f"Error: Cannot submit in '{self._state}' state.")
+    def __init__(self):
+        # The initial state is always Draft
+        self.transition_to(DraftState())
 
+    def transition_to(self, state: State):
+        print(f"Transitioning to {state.name}...")
+        self._state = state
+        self._state.set_context(self)
+
+    # The context delegates all work to the current state
     def publish(self):
-        if self._state == "under_review":
-            self._state = "published"
-            print(f"Document published. New state: '{self._state}'")
-        else:
-            print(f"Error: Cannot publish from '{self._state}' state.")
+        self._state.publish()
 
     def reject(self):
-        if self._state == "under_review":
-            # In a real app, you might have more complex logic here
-            self._state = "needs_revision"
-            print(f"Document needs revision. New state: '{self._state}'")
-        else:
-            print(f"Error: Cannot reject from '{self._state}' state.")
+        self._state.reject()
+
+    def edit(self):
+        self._state.edit()
+
+    def submit(self):
+        self._state.submit()

The refactored document.py is clean, concise, and follows the Single Responsibility Principle. Its only job is to manage the current state and delegate tasks.

Visualizing the Complete Workflow

With our new structure, we can clearly visualize the entire document lifecycle as a state machine.

stateDiagram-v2
    [*] --> Draft
    Draft --> UnderReview: submit()
    UnderReview --> Published: publish()
    UnderReview --> NeedsRevision: reject()
    NeedsRevision --> UnderReview: submit()
    Published --> [*]

Step 5: Putting It All Together

Our main.py file will simulate user actions and drive the state transitions.

# main.py
from document import Document

def main():
    doc = Document()  # Starts in DraftState

    print("\n--- Attempting to publish from Draft (should fail) ---")
    doc.publish()

    print("\n--- Editing and Submitting ---")
    doc.edit()
    doc.submit()  # Moves to UnderReviewState

    print("\n--- Attempting to edit while under review (should fail) ---")
    doc.edit()

    print("\n--- Rejecting the document ---")
    doc.reject()  # Moves to NeedsRevisionState

    print("\n--- Editing and re-submitting ---")
    doc.edit()
    doc.submit()  # Moves back to UnderReviewState

    print("\n--- Publishing the document ---")
    doc.publish()  # Moves to PublishedState

    print("\n--- Attempting to edit after publishing (should fail) ---")
    doc.edit()

if __name__ == "__main__":
    main()

Benefits of the State Pattern

1. Single Responsibility Principle (SRP): Each state class has one job: to implement the logic for that state.
2. Open/Closed Principle (OCP): To add a new state (e.g., ArchivedState), you create a new class. You don’t need to modify existing state classes or the context.
3. Improved Readability: The logic is no longer a monolithic block of conditionals. It’s organized into clean, self-contained classes.
4. Simplified Context: The Document class is simple and stable. It doesn’t care about the specific state logic, only that it has a state to delegate to.

Conclusion

The State Design Pattern is a powerful tool for managing complex, state-driven behavior in an object-oriented way. By refactoring our initial if/elif mess into a collection of state objects, we’ve created a system that is not only easier to understand and maintain but also flexible enough to grow without collapsing under its own complexity.