python - Implementing a state machine with decorators -

while learning concepts of decorators in python came question if possible use decorators simulate state machine.

example:

from enum import enum   class coffeemachine(object):     def __init__(self):         self.state = coffeestate.initial      #@statemachine(shouldbe, willbe)     @statemachine(coffeestate.initial, coffeestate.grounding)     def ground_beans(self):         print("ground_beans")      @statemachine(coffeestate.grounding, coffeestate.heating)     def heat_water(self):         print("heat_water")      @statemachine(coffeestate.heating, coffeestate.pumping)     def pump_water(self):         print("pump_water")   class coffeestate(enum):     initial = 0     grounding = 1     heating = 2     pumping = 3 

so statemachine check if current state requested one, if is, should call underlying function , lastly should set state further.

how implement this?

sure can, provided decorator makes assumption state stored:

from functools import wraps   class statemachinewrongstate(exception):     def __init__(self, shouldbe, current):         self.shouldbe = shouldbe         self.current = current         super().__init__((shouldbe, current))   def statemachine(shouldbe, willbe):     def decorator(f):         @wraps(f)         def wrapper(self, *args, **kw):             if self.state != shouldbe:                 raise statemachinewrongstate(shouldbe, self.state)             try:                 return f(self, *args, **kw)             finally:                 self.state = willbe         return wrapper     return decorator 

the decorator expects self passed in; i.e. should applied methods in class. expects self have state attribute track state machine state.

demo:

>>> cm = coffeemachine() >>> cm.state <coffeestate.initial: 0> >>> cm.ground_beans() ground_beans >>> cm.state <coffeestate.grounding: 1> >>> cm.ground_beans() traceback (most recent call last):   file "<stdin>", line 1, in <module>   file "<stdin>", line 6, in wrapper __main__.statemachinewrongstate: (<coffeestate.initial: 0>, <coffeestate.grounding: 1>) >>> cm.heat_water() heat_water >>> cm.pump_water() pump_water >>> cm.state <coffeestate.pumping: 3>