Data/ML Engineer Blog

State (behavior changes with internal state)

When to use

  • An object’s allowed actions depend on its state (queued → running → succeeded/failed).
  • You want to avoid big if state == ... ladders scattered around.
  • You need explicit, safe transitions with good errors for illegal moves.

Avoid when states are simple—an Enum + small match might be enough.

Diagram (text)

PipelineRun (context) ─ holds ─> RunState
   start() ──> Running
   succeed() ──> Succeeded
   fail() ──> Failed(reason)
Queued ── start() → Running
Running ── succeed() → Succeeded; fail() → Failed
Succeeded/Failed ── all actions illegal

Python example (≤40 lines, type-hinted)

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Protocol

class RunState(Protocol):
    def start(self, ctx: "PipelineRun") -> None: ...
    def succeed(self, ctx: "PipelineRun") -> None: ...
    def fail(self, ctx: "PipelineRun", reason: str) -> None: ...

def _illegal(action: str) -> None:
    raise RuntimeError(f"illegal transition: {action}")

class Queued:
    def start(self, ctx: "PipelineRun") -> None: ctx.state = Running()
    def succeed(self, ctx: "PipelineRun") -> None: _illegal("queued.succeed")
    def fail(self, ctx: "PipelineRun", reason: str) -> None: ctx.state = Failed(reason)

class Running:
    def start(self, ctx: "PipelineRun") -> None: _illegal("running.start")
    def succeed(self, ctx: "PipelineRun") -> None: ctx.state = Succeeded()
    def fail(self, ctx: "PipelineRun", reason: str) -> None: ctx.state = Failed(reason)

class Succeeded:
    def start(self, ctx: "PipelineRun") -> None: _illegal("succeeded.start")
    def succeed(self, ctx: "PipelineRun") -> None: _illegal("succeeded.succeed")
    def fail(self, ctx: "PipelineRun", reason: str) -> None: _illegal("succeeded.fail")

class Failed(Succeeded):
    def __init__(self, reason: str) -> None: self.reason = reason
    def fail(self, ctx: "PipelineRun", reason: str) -> None: _illegal("failed.fail")

@dataclass
class PipelineRun:
    state: RunState = field(default_factory=Queued)
    def start(self) -> None: self.state.start(self)
    def succeed(self) -> None: self.state.succeed(self)
    def fail(self, reason: str) -> None: self.state.fail(self, reason)

Tiny pytest (cements it)

def test_state_transitions_and_guards():
    run = PipelineRun()
    assert isinstance(run.state, Queued)
    run.start(); assert isinstance(run.state, Running)
    run.succeed(); assert isinstance(run.state, Succeeded)

def test_illegal_and_failure_reason():
    run = PipelineRun(); run.start(); run.fail("timeout")
    assert isinstance(run.state, Failed) and run.state.reason == "timeout"
    import pytest; with pytest.raises(RuntimeError): run.start()

Trade-offs & pitfalls

  • Pros: Localizes rules; removes if/elif chains; illegal transitions fail fast; easy to add states.
  • Cons: Extra small classes; some indirection when reading the flow.
  • Pitfalls:
    • Letting states carry mutable shared data—keep data on the context (PipelineRun), not the states.
    • Hidden transitions—document them; tests should cover the graph.
    • Overusing State when an Enum + match would be simpler.

Pythonic alternatives

  • Enum + match/dict of transitions for simple graphs.
  • Dataclass context + Strategy if you only vary behavior, not transitions.
  • Finite State Machine libs (e.g., transitions) if you need guards/entry/exit actions/timers.

Mini exercise

Add a cancel() action allowed from Queued and Running, leading to a new Canceled state (other actions become illegal there). Write tests for valid cancels and illegal ones (e.g., cancel after success).

Checks (quick checklist)

  • Context holds current state object; methods delegate to it.
  • Each state implements allowed actions and raises on illegal ones.
  • Data lives on the context, not in state singletons.
  • Tests cover all valid/invalid transitions.
  • Prefer Enum + match if the graph is tiny.