使用装饰器实现状态机

Implementing a state machine with decorators

在 python 中学习装饰器的概念时,我想到了是否可以使用装饰器来模拟状态机的问题。

示例:

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

所以状态机所做的就是检查我的当前状态是否是请求的状态,如果是,它应该调用底层函数,最后它应该进一步设置状态。

您将如何实施?

当然可以,前提是您的装饰器对状态的存储位置做出了假设:

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

装饰器期望得到self传入;即它应该应用于 class 中的方法。然后它期望 self 有一个 state 属性来跟踪状态机状态。

演示:

>>> 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>