Python State Pattern: A Visual Guide to Taming Complex Workflows

Have you ever found yourself trapped in a labyrinth of if/elif/else statements, where an object’s behavior changes based on its current “state”? This common scenario can quickly lead to code that’s difficult to read, maintain, and extend.

Today, we’ll explore a powerful solution: the State Design Pattern. We’ll start with a familiar problem—a document workflow system—and refactor it from a complex, monolithic class into a clean, scalable, and state-driven machine.

The Initial Problem: A Monolithic Document Workflow

Imagine a system for managing articles. An article, or Document, moves through several stages before it’s published.

Here’s a visual representation of the workflow:

graph TD;
    A[Draft] -->|submit()| B(Under Review);
    B -->|approve()| C{Published};
    B -->|request_changes()| D(Needs Revision);
    D -->|submit()| B;
    B -->|reject()| E[Rejected];

A straightforward approach might be to manage these states with a single property (e.g., self.state) inside the Document class and use conditional logic in each method.

Let’s see what that looks like in code.

The “Before” Code

class Document:
    def __init__(self, content: str):
        self.content = content
        self.state = "draft"  # Initial state
        print(f"Document created. Initial state: '{self.state}'")

    def edit(self, new_content: str):
        print("Attempting to edit document...")
        if self.state in ["draft", "needs_revision"]:
            self.content = new_content
            print("Document edited successfully.")
        elif self.state == "under_review":
            print("[ERROR] Cannot edit while under review.")
        elif self.state == "published":
            print("[ERROR] Cannot edit a published document.")
        elif self.state == "rejected":
            print("[ERROR] Cannot edit a rejected document.")
        else:
            print(f"[ERROR] Unknown state: {self.state}")

    def submit(self):
        print("Attempting to submit document...")
        if self.state in ["draft", "needs_revision"]:
            self.state = "under_review"
            print(f"Document submitted. New state: '{self.state}'")
        elif self.state == "under_review":
            print("[ERROR] Document is already under review.")
        else:
            print(f"[ERROR] Cannot submit from state: '{self.state}'")

    def publish(self):
        print("Attempting to publish document...")
        if self.state == "under_review":
            self.state = "published"
            print(f"Document published! New state: '{self.state}'")
        elif self.state == "draft":
            print("[ERROR] Cannot publish directly from draft. Please submit for review first.")
        else:
            print(f"[ERROR] Cannot publish from state: '{self.state}'")

    def reject(self, final: bool = False):
        print("Attempting to reject or request revisions...")
        if self.state == "under_review":
            if final:
                self.state = "rejected"
                print(f"Document rejected. New state: '{self.state}'")
            else:
                self.state = "needs_revision"
                print(f"Revisions requested. New state: '{self.state}'")
        else:
            print(f"[ERROR] Cannot reject from state: '{self.state}'")

This works, but it’s already showing signs of trouble. Each method is a tangled web of conditionals.

The Breaking Point: Adding a New State

Now, the business asks for a new feature: documents can become expired if they stay in under_review or needs_revision for too long. Once expired, no further actions should be possible.

To implement this, we’d have to modify every single method in the Document class.

Here’s how the edit method would change:

--- a/document.py
+++ b/document.py
@@ -6,6 +6,8 @@
     def edit(self, new_content: str):
         print("Attempting to edit document...")
         if self.state in ["draft", "needs_revision"]:
             self.content = new_content
             print("Document edited successfully.")
+        elif self.state == "expired":
+            print("[ERROR] Cannot edit an expired document.")
         elif self.state == "under_review":
             print("[ERROR] Cannot edit while under review.")
         elif self.state == "published":

This is a classic violation of the Open/Closed Principle. Our class should be open for extension but closed for modification. Every new state forces us to hunt down and modify existing, working code, which is a recipe for bugs.

[!NOTE] The State Pattern is a behavioral design pattern that allows an object to change its behavior when its internal state changes. The object appears to have changed its class.

The Solution: Introducing the State Pattern

The State pattern elegantly solves this by encapsulating the logic for each state into its own separate class. The main object, called the Context, holds a reference to a state object and delegates all state-specific work to it.

Here is the UML diagram representing our new architecture:

classDiagram
    class Document {
        -state: State
        +setState(state)
        +edit(content)
        +submit()
        +publish()
    }
    class State {
        <<interface>>
        +document: Document
        +edit(content)
        +submit()
        +publish()
    }
    Document o-- State
    State <|-- DraftState
    State <|-- UnderReviewState
    State <|-- PublishedState
    State <|-- NeedsRevisionState

    class DraftState {
        +submit()
    }
    class UnderReviewState {
        +publish()
        +reject()
    }
    class PublishedState {
        // No actions allowed
    }
    class NeedsRevisionState {
        +edit(content)
        +submit()
    }

Let’s refactor our code following this structure.

Step 1: Define the State Interface

We’ll create an abstract base class (ABC) that defines the common interface for all states. By default, each action will raise an error, forcing concrete states to implement only the actions they support.

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

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.
    """
    _document: Document

    @property
    def document(self) -> Document:
        return self._document

    @document.setter
    def document(self, document: Document) -> None:
        self._document = document

    def edit(self, new_content: str) -> None:
        print(f"[ERROR] Cannot edit in '{self.__class__.__name__}' state.")

    def submit(self) -> None:
        print(f"[ERROR] Cannot submit in '{self.__class__.__name__}' state.")

    def publish(self) -> None:
        print(f"[ERROR] Cannot publish in '{self.__class__.__name__}' state.")

    def reject(self, final: bool) -> None:
        print(f"[ERROR] Cannot reject in '{self.__class__.__name__}' state.")

    def expire(self) -> None:
        print(f"[ERROR] Cannot expire in '{self.__class__.__name__}' state.")

Step 2: Create Concrete State Classes

Now, we create a class for each state in our workflow. Each class inherits from State and implements only the methods relevant to it.

Step 3: Refactor the Context (Document) Class

The Document class becomes dramatically simpler. It holds a state object, delegates calls to it, and provides a method to transition to a new state.

--- a/document_before.py
+++ b/document_after.py
@@ -1,45 +1,33 @@
-class Document:
-    def __init__(self, content: str):
-        self.content = content
-        self.state = "draft"
-        print(f"Document created. Initial state: '{self.state}'")
+from states.base_state import State
+from states.draft_state import DraftState
 
-    def edit(self, new_content: str):
-        print("Attempting to edit document...")
-        if self.state in ["draft", "needs_revision"]:
-            self.content = new_content
-            print("Document edited successfully.")
-        elif self.state == "under_review":
-            print("[ERROR] Cannot edit while under review.")
-        elif self.state == "published":
-            print("[ERROR] Cannot edit a published document.")
-        elif self.state == "rejected":
-            print("[ERROR] Cannot edit a rejected document.")
-        else:
-            print(f"[ERROR] Unknown state: {self.state}")
+class Document:
+    """
+    The Context defines the interface of interest to clients. It also
+    maintains a reference to an instance of a State subclass, which
+    represents the current state of the Context.
+    """
+    _state: State = None
+
+    def __init__(self, content: str):
+        self.content = content
+        self.transition_to(DraftState())
+
+    def transition_to(self, state: State):
+        print(f"Transitioning to {type(state).__name__}...")
+        self._state = state
+        self._state.document = self
 
     def submit(self):
-        print("Attempting to submit document...")
-        if self.state in ["draft", "needs_revision"]:
-            self.state = "under_review"
-            print(f"Document submitted. New state: '{self.state}'")
-        elif self.state == "under_review":
-            print("[ERROR] Document is already under review.")
-        else:
-            print(f"[ERROR] Cannot submit from state: '{self.state}'")
+        self._state.submit()
+
+    def edit(self, new_content: str):
+        self._state.edit(new_content)
 
     def publish(self):
-        print("Attempting to publish document...")
-        if self.state == "under_review":
-            self.state = "published"
-            print(f"Document published! New state: '{self.state}'")
-        elif self.state == "draft":
-            print("[ERROR] Cannot publish directly from draft. Please submit for review first.")
-        else:
-            print(f"[ERROR] Cannot publish from state: '{self.state}'")
+        self._state.publish()
 
     def reject(self, final: bool = False):
-        print("Attempting to reject or request revisions...")
-        if self.state == "under_review":
-            if final:
-                self.state = "rejected"
-                print(f"Document rejected. New state: '{self.state}'")
-            else:
-                self.state = "needs_revision"
-                print(f"Revisions requested. New state: '{self.state}'")
-        else:
-            print(f"[ERROR] Cannot reject from state: '{self.state}'")
+        self._state.reject(final)
+
+    def expire(self):
+        self._state.expire()

Look at how clean that is! The Document class no longer contains any conditional logic about its state. It simply delegates the responsibility to its current state object. Adding a new state now means creating a new class, not modifying a dozen methods in an existing one.

[!TIP] State vs. Strategy: The State and Strategy patterns have very similar structures but different intents.

• State: Manages an object’s behavior as its internal state changes. The states are often instantiated and managed by the context object itself.
• Strategy: Provides a family of interchangeable algorithms for a specific task. The client usually provides the context with the desired strategy object.

Summary and Key Benefits

By applying the State pattern, we’ve achieved several key benefits:

mindmap
  root((State Pattern Benefits))
    **Single Responsibility Principle**
      Each state class has one job: manage the logic for that specific state.
    **Open/Closed Principle**
      We can add new states without modifying existing state classes or the context class.
    **Clean Code**
      Replaced complex conditional blocks with clean, polymorphic state objects.
    **Improved Cohesion**
      State-specific code is now grouped together in its own class, rather than being scattered across methods in the context.

The State pattern is an indispensable tool for any developer dealing with objects that have complex, state-dependent behavior. It allows you to transform a tangled mess of conditionals into a clean, maintainable, and robust state machine that is a pleasure to work with.