Python 如何使用装饰器实现有限状态机？

Z时代
2024-02-09
分类：IT

我有一个需求，需要类似有限状态机的功能，比如一个工作流程，有 4 个步骤，没个步骤都可能引发某种异常，或者返回某个值，这导致下一步需要执行的步骤存在多种可能，常规的写法是一大坨 if-else，另外，步骤之间可能需要回滚重试，这使得 if-else 也不太好写。
我构想了一种方式，灵感来自有限状态机，初步的期望用法如下，但是我不知道要如何去实现？

@fsm(condition=Exception, method='step2')
@fsm(condition=True, method='step3')
@fsm(condition=False, method='step4')
def step1(arg):
    return arg
@fsm(condition=Exception, method='step2', retry=3)
@fsm(condition=True, method='step3')
@fsm(condition=False, method='step4')
def step2(arg):
    return arg
@fsm(condition=Exception, method='step1', retry=3)
@fsm(condition=True, method='step3')
@fsm(condition=False, method='step4')
def step3(arg):
    return arg
@fsm(condition=Exception, retry=3)
def step4(arg):
    return arg
# 启动流程。
step1('arg')

回答：

推测一下你想要实现的功能：

转发。实现一个函数装饰器，参数两个：条件值(condition)和后续动作函数(method)，当当前函数运行结束并返回的时，判断返回值是否与条件值匹配，如果匹配，调用后续动作函数。
重试。实现一个函数装饰圈，参数两个：重试次数(retry)和全部失败后的后续动作函数(method)。如果当前函数运行失败，且失败次数小于重试次数，用原参数重试当前函数；如果超过重试次数，调用后续动作函数收尾。

关于转发和重试时，后续动作函数的参数 arg，从你的描述中不太明白该取何值，将原 arg 传递到下游。代码实现如下：

python">def fsmError(retry, method):
    def newFunc(f):
        def g(args):
            try:
                return f(args)
            except:
                if retry > 0:
                    return fsmError(retry - 1, method)(f)(args)
                else:
                    d = globals()[method]
                    return d(args)
        return g
    return newFunc
def fsm(condition, method):
    def newFunc(f):
        def g(args):
            res = f(args)
            if condition == res:
                n = globals()[method]
                return n(args)
            return res
        return g
    return newFunc
mockRetry = 0
@fsmError(retry=5, method="error")
@fsm(condition=True, method="step2")
@fsm(condition=False, method="step3")
def step1(args):
    global mockRetry
    mockRetry += 1
    print("第" + str(mockRetry) + "次调用")
    if mockRetry < 3:
        1 / 0
    print("step1")
    return True
@fsm(condition=None, method="step3")
def step2(args):
    print("step2")
def step3(args):
    print("step3")
def error(args):
    print("重试失败")
if __name__ == "__main__":
    res = step1(0)