The State Pattern in Python: Refactoring a Vending Machine Into State Classes

Lesson 2 built a vending machine whose methods each guard on the current state: select_item refuses anything but HAS_MONEY, dispense anything but DISPENSING. That works until a new state arrives whose transitions cut across the others.

Adding a state to the if/elif machine

The operator wants a MAINTENANCE state. Here is lesson 2’s code unchanged, plus the two methods for the technician’s buttons.

# fsm_maintenance_bug.py: lesson 2's machine, plus the technician's two maintenance buttons
class VendingMachine:
    def __init__(self):
        self.state = "IDLE"
        self._stock = 1

    def insert_coin(self):
        if self.state == "IDLE":
            self.state = "HAS_MONEY"
            print("Coin accepted. State -> HAS_MONEY")
        else:
            print(f"Cannot insert coin in state {self.state}")

    def select_item(self):
        if self.state != "HAS_MONEY":
            print(f"Cannot select item in state {self.state}: insert coin first")
            return
        if self._stock <= 0:
            self.state = "SOLD_OUT"
            print("Out of stock. State -> SOLD_OUT")
            return
        self.state = "DISPENSING"
        print("Item selected. State -> DISPENSING")

    def dispense(self):
        if self.state != "DISPENSING":
            print(f"Cannot dispense in state {self.state}")
            return
        self._stock -= 1
        self.state = "IDLE"
        print(f"Dispensed. Stock left: {self._stock}. State -> IDLE")

    def start_maintenance(self):              # new
        print(f"Maintenance started (was {self.state}). State -> MAINTENANCE")
        self.state = "MAINTENANCE"

    def end_maintenance(self):                # new
        self.state = "IDLE"
        print("Maintenance done. State -> IDLE")

m = VendingMachine()
m.insert_coin()
m.start_maintenance()                         # a customer's coin is inside
m.select_item()
m.end_maintenance()
m.select_item()

$ python fsm_maintenance_bug.py
Coin accepted. State -> HAS_MONEY
Maintenance started (was HAS_MONEY). State -> MAINTENANCE
Cannot select item in state MAINTENANCE: insert coin first
Maintenance done. State -> IDLE
Cannot select item in state IDLE: insert coin first

The guards did their job: select_item refused correctly while under maintenance. The bug is in the new transition. The customer’s coin was accepted, maintenance started, and when it ended the machine was back in IDLE asking for a coin. Nothing refunded the first one. start_maintenance jumps from any state, but it doesn’t know that HAS_MONEY means “holding a coin”. Fixing it inside that method means an if for HAS_MONEY, another for DISPENSING, and one more each time a state is added: the rules for each state end up spread across methods again.

State pattern: one class per state

Each state becomes a class with the same events. An event returns the next state, and the machine only forwards events to whatever state it’s in.

# state_classes.py: lesson 2's states as classes; each event returns the next state
class State:
    """Default: every event is refused. A state overrides only what it allows."""
    name = "?"
    def insert_coin(self, m):
        print(f"Cannot insert coin in state {self.name}")
        return self
    def select_item(self, m):
        print(f"Cannot select item in state {self.name}")
        return self
    def dispense(self, m):
        print(f"Cannot dispense in state {self.name}")
        return self
    def start_maintenance(self, m):
        print(f"Maintenance started (was {self.name})")
        return Maintenance()
    def end_maintenance(self, m):
        print(f"Not under maintenance ({self.name})")
        return self

class Idle(State):
    name = "IDLE"
    def insert_coin(self, m):
        print("Coin accepted")
        return HasMoney()

class HasMoney(State):
    name = "HAS_MONEY"
    def select_item(self, m):
        if m.stock <= 0:
            print("Out of stock, coin returned")
            return SoldOut()
        print("Item selected")
        return Dispensing()
    def start_maintenance(self, m):           # this state knows it holds a coin
        print("Coin returned to the customer")
        return super().start_maintenance(m)

class Dispensing(State):
    name = "DISPENSING"
    def dispense(self, m):
        m.stock -= 1
        print(f"Dispensed. Stock left: {m.stock}")
        return Idle()

class SoldOut(State):
    name = "SOLD_OUT"

class Maintenance(State):                     # the new state is one new class
    name = "MAINTENANCE"
    def start_maintenance(self, m):
        print("Already under maintenance")
        return self
    def end_maintenance(self, m):
        print("Maintenance done")
        return Idle() if m.stock > 0 else SoldOut()

class VendingMachine:
    """Holds the data and the current state; each event asks the state what happens next."""
    def __init__(self):
        self.state, self.stock = Idle(), 1
    def insert_coin(self):
        self.state = self.state.insert_coin(self)
    def select_item(self):
        self.state = self.state.select_item(self)
    def dispense(self):
        self.state = self.state.dispense(self)
    def start_maintenance(self):
        self.state = self.state.start_maintenance(self)
    def end_maintenance(self):
        self.state = self.state.end_maintenance(self)

m = VendingMachine()
m.insert_coin()
m.start_maintenance()
m.select_item()
m.end_maintenance()
m.insert_coin()                               # the customer tries again
m.select_item()
m.dispense()
print("state:", m.state.name)

$ python state_classes.py
Coin accepted
Coin returned to the customer
Maintenance started (was HAS_MONEY)
Cannot select item in state MAINTENANCE
Maintenance done
Coin accepted
Item selected
Dispensed. Stock left: 0
state: IDLE

The same sequence now refunds the coin, because HasMoney overrides start_maintenance and then calls the shared version. The state that holds the coin is the one that knows to return it. Selecting under maintenance is still refused, and afterwards a fresh purchase works through HAS_MONEY and DISPENSING as in lesson 2.

Two design choices carry this:

• The base class refuses by default. Maintenance only overrides the two maintenance events; inserting, selecting and dispensing are refused because State already refuses them. A new state is safe until its author allows something.
• All of one state’s rules live in one class. Adding MAINTENANCE was one new class plus one override in HasMoney. No other method needed reading, so none could be missed.

The cost is more classes for a small machine. With four states and three events, lesson 2’s version is shorter and perfectly fine; the State pattern pays off when states keep being added or have behaviour of their own, like the refund.

Practice

# p17_1.py: maintenance must wait while an item is dropping, so DISPENSING refuses it
class State:
    name = "?"
    def select_item(self, m):
        print(f"Cannot select item in state {self.name}")
        return self
    def dispense(self, m):
        print(f"Cannot dispense in state {self.name}")
        return self
    def start_maintenance(self, m):
        print(f"Maintenance started (was {self.name})")
        return Maintenance()

class HasMoney(State):
    name = "HAS_MONEY"
    def select_item(self, m):
        print("Item selected")
        return Dispensing()

class Dispensing(State):
    name = "DISPENSING"
    def dispense(self, m):
        m.stock -= 1
        print(f"Dispensed. Stock left: {m.stock}")
        return Idle()
    def start_maintenance(self, m):           # the one new override
        print("Wait: an item is dropping")
        return self

class Idle(State):
    name = "IDLE"

class Maintenance(State):
    name = "MAINTENANCE"

class VendingMachine:
    def __init__(self):
        self.state, self.stock = HasMoney(), 1
    def select_item(self):
        self.state = self.state.select_item(self)
    def dispense(self):
        self.state = self.state.dispense(self)
    def start_maintenance(self):
        self.state = self.state.start_maintenance(self)

m = VendingMachine()
m.select_item()
m.start_maintenance()
m.dispense()
m.start_maintenance()
print("state:", m.state.name)

$ python p17_1.py
Item selected
Wait: an item is dropping
Dispensed. Stock left: 0
Maintenance started (was IDLE)
state: MAINTENANCE

# p17_2.py: an item costs two coins, so HAS_MONEY keeps a balance until it's enough
PRICE = 2

class State:
    name = "?"
    def insert_coin(self, m):
        print(f"Cannot insert coin while {self.name}")
        return self
    def select_item(self, m):
        print(f"Cannot select while {self.name}")
        return self

class Idle(State):
    name = "IDLE"
    def insert_coin(self, m):
        m.balance = 1
        print("Coin accepted, balance 1")
        return HasMoney()

class HasMoney(State):
    name = "HAS_MONEY"
    def insert_coin(self, m):
        m.balance += 1
        print(f"Coin accepted, balance {m.balance}")
        return self
    def select_item(self, m):
        if m.balance < PRICE:
            print(f"Need {PRICE - m.balance} more coin")
            return self
        m.balance -= PRICE
        print(f"Dispensed, change returned: {m.balance}")
        m.balance = 0
        return Idle()

class VendingMachine:
    def __init__(self):
        self.state, self.balance = Idle(), 0
    def insert_coin(self):
        self.state = self.state.insert_coin(self)
    def select_item(self):
        self.state = self.state.select_item(self)

m = VendingMachine()
m.insert_coin()
m.select_item()
m.insert_coin()
m.insert_coin()
m.select_item()
print("state:", m.state.name)

$ python p17_2.py
Coin accepted, balance 1
Need 1 more coin
Coin accepted, balance 2
Coin accepted, balance 3
Dispensed, change returned: 1
state: IDLE

# p17_3.py: restock() is allowed only under maintenance; every other state refuses by default
class State:
    name = "?"
    def restock(self, m, units):
        print(f"Cannot restock while {self.name}")
        return self
    def start_maintenance(self, m):
        print("Technician opened the machine")
        return Maintenance()

class SoldOut(State):
    name = "SOLD_OUT"

class Maintenance(State):
    name = "MAINTENANCE"
    def restock(self, m, units):
        m.stock += units
        print(f"Restocked {units}, stock now {m.stock}")
        return self

class VendingMachine:
    def __init__(self):
        self.state, self.stock = SoldOut(), 0
    def restock(self, units):
        self.state = self.state.restock(self, units)
    def start_maintenance(self):
        self.state = self.state.start_maintenance(self)

m = VendingMachine()
m.restock(10)
m.start_maintenance()
m.restock(10)
print("stock:", m.stock)

$ python p17_3.py
Cannot restock while SOLD_OUT
Technician opened the machine
Restocked 10, stock now 10
stock: 10

# p17_4.py: the machine logs every transition, because state changes happen in one place
class State:
    name = "?"
    def insert_coin(self, m):
        return self
    def select_item(self, m):
        return self

class Idle(State):
    name = "IDLE"
    def insert_coin(self, m):
        return HasMoney()

class HasMoney(State):
    name = "HAS_MONEY"
    def select_item(self, m):
        m.stock -= 1
        return Idle() if m.stock > 0 else SoldOut()

class SoldOut(State):
    name = "SOLD_OUT"

class VendingMachine:
    def __init__(self, stock):
        self.state, self.stock, self.log = Idle(), stock, []
    def _fire(self, event):
        before = self.state.name
        self.state = getattr(self.state, event)(self)
        self.log.append(f"{event}: {before} -> {self.state.name}")
    def insert_coin(self):
        self._fire("insert_coin")
    def select_item(self):
        self._fire("select_item")

m = VendingMachine(stock=1)
m.select_item()
m.insert_coin()
m.select_item()
m.insert_coin()
print("\n".join(m.log))

$ python p17_4.py
select_item: IDLE -> IDLE
insert_coin: IDLE -> HAS_MONEY
select_item: HAS_MONEY -> SOLD_OUT
insert_coin: SOLD_OUT -> SOLD_OUT