Python生成器与状态机实现
Python生成器与状态机实现
生成器可以看作是一个保存了执行状态的函数。每次yield暂停执行并保存状态，下次调用send恢复执行。这个特性恰好可以用来实现状态机。
一个典型的状态机实现：
import functools
def state_machine(initial_state):
def decorator(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
gen = func(*args, **kwargs)
# 推进到第一个yield
next(gen)
return gen
return wrapper
return decorator
def transition(event):
def decorator(func):
func._is_transition = True
func._event = event
return func
return decorator
class StateMachine:
def __init__(self):
self.state = None
self._transitions = {}
self._register_transitions()
def _register_transitions(self):
for attr_name in dir(self):
attr = getattr(self, attr_name)
if hasattr(attr, '_is_transition') and attr._is_transition:
event = attr._event
target_state = attr(self)
self._transitions[(self.__class__.__name__, event)] = target_state
@state_machine('idle')
def tcp_connection():
state = 'closed'
while True:
event = yield
if state == 'closed':
if event == 'connect':
print("CLOSED -> SYN_SENT: 发送SYN")
state = 'syn_sent'
elif event == 'passive_open':
print("CLOSED -> LISTEN: 开始监听")
state = 'listen'
elif state == 'syn_sent':
if event == 'syn_ack':
print("SYN_SENT -> ESTABLISHED: 连接建立")
state = 'established'
elif event == 'timeout':
print("SYN_SENT -> CLOSED: 超时关闭")
state = 'closed'
elif state == 'established':
if event == 'fin':
print("ESTABLISHED -> FIN_WAIT_1: 主动关闭")
state = 'fin_wait_1'
elif event == 'data':
print("ESTABLISHED: 接收数据")
elif state == 'fin_wait_1':
if event == 'ack':
print("FIN_WAIT_1 -> FIN_WAIT_2: 等待对方关闭")
state = 'fin_wait_2'
elif event == 'fin':
print("FIN_WAIT_1 -> CLOSING: 同时关闭")
state = 'closing'
elif state == 'fin_wait_2':
if event == 'fin':
print("FIN_WAIT_2 -> TIME_WAIT: 收到对方关闭请求")
state = 'time_wait'
elif state == 'time_wait':
if event == 'timeout':
print("TIME_WAIT -> CLOSED: 关闭")
state = 'closed'
conn = tcp_connection()
conn.send('connect') # CLOSED -> SYN_SENT
conn.send('syn_ack') # SYN_SENT -> ESTABLISHED
conn.send('data') # 接收数据
conn.send('fin') # ESTABLISHED -> FIN_WAIT_1
生成器状态机的好处是状态被保存在生成器的局部变量中，不需要额外的类属性。但缺点是状态转换逻辑集中在一个函数中，随着状态增多变得难以维护。
更好的方式是用类组织状态，用生成器组织每个状态内的逻辑：
class ConnectionFSM:
def __init__(self):
self.state = 'closed'
self._buffer = []
def closed(self, event):
if event == 'connect':
print("CLOSED -> SYN_SENT")
self.state = 'syn_sent'
return True
elif event == 'passive_open':
print("CLOSED -> LISTEN")
self.state = 'listen'
return True
return False
def syn_sent(self, event):
if event == 'syn_ack':
print("SYN_SENT -> ESTABLISHED")
self.state = 'established'
return True
elif event == 'timeout':
print("SYN_SENT -> CLOSED (timeout)")
self.state = 'closed'
return True
return False
def dispatch(self, event):
handler = getattr(self, self.state, None)
if handler is None:
raise ValueError(f"Unknown state: {self.state}")
return handler(event)
fsm = ConnectionFSM()
fsm.dispatch('connect')
fsm.dispatch('syn_ack')
fsm.dispatch('data')
每个状态是一个独立的方法，状态转换通过修改self.state实现。dispatch方法根据当前状态自动路由事件。
用枚举定义状态表更清晰：
from enum import Enum, auto
class TCPState(Enum):
CLOSED = auto()
LISTEN = auto()
SYN_SENT = auto()
SYN_RCVD = auto()
ESTABLISHED = auto()
FIN_WAIT_1 = auto()
FIN_WAIT_2 = auto()
CLOSING = auto()
TIME_WAIT = auto()
class TCPConnection:
def __init__(self):
self.state = TCPState.CLOSED
self._transitions = {
(TCPState.CLOSED, 'connect'): TCPState.SYN_SENT,
(TCPState.CLOSED, 'passive_open'): TCPState.LISTEN,
(TCPState.SYN_SENT, 'syn_ack'): TCPState.ESTABLISHED,
(TCPState.SYN_SENT, 'timeout'): TCPState.CLOSED,
(TCPState.ESTABLISHED, 'fin'): TCPState.FIN_WAIT_1,
(TCPState.ESTABLISHED, 'data'): TCPState.ESTABLISHED,
(TCPState.FIN_WAIT_1, 'ack'): TCPState.FIN_WAIT_2,
(TCPState.FIN_WAIT_1, 'fin'): TCPState.CLOSING,
(TCPState.FIN_WAIT_2, 'fin'): TCPState.TIME_WAIT,
(TCPState.CLOSING, 'ack'): TCPState.TIME_WAIT,
(TCPState.TIME_WAIT, 'timeout'): TCPState.CLOSED,
}
def send(self, event):
key = (self.state, event)
if key in self._transitions:
old_state = self.state
self.state = self._transitions[key]
self._on_transition(old_state, self.state, event)
else:
raise ValueError(f"Invalid transition: {self.state} -> {event}")
def _on_transition(self, from_state, to_state, event):
print(f"{from_state.name} --({event})--> {to_state.name}")
状态转换表用字典集中管理，每个状态转移可以附加进入动作、退出动作和转换动作：
class StateMachineMeta:
def __init__(self):
self._states = {}
self._current = None
def add_state(self, name, on_enter=None, on_exit=None):
self._states[name] = {
'on_enter': on_enter,
'on_exit': on_exit,
'transitions': {}
}
def add_transition(self, from_state, event, to_state, action=None):
self._states[from_state]['transitions'][event] = (to_state, action)
def send(self, event):
current_state = self._states[self._current]
if event not in current_state['transitions']:
raise ValueError(f"No transition for {event} from {self._current}")
to_state, action = current_state['transitions'][event]
if current_state.get('on_exit'):
current_state['on_exit'](self._current)
if action:
action()
if self._states[to_state].get('on_enter'):
self._states[to_state]['on_enter'](to_state)
self._current = to_state
sm = StateMachineMeta()
sm.add_state('idle', on_enter=lambda s: print(f"进入{s}"))
sm.add_state('running', on_enter=lambda s: print(f"进入{s}"))
sm.add_transition('idle', 'start', 'running', action=lambda: print("启动..."))
生成器状态机的一个优势是可以方便地实现嵌套状态机：
def nested_fsm():
outer_state = 'active'
inner_state = 'idle'
while True:
event = yield
if outer_state == 'active':
if inner_state == 'idle':
if event == 'start':
inner_state = 'working'
print("开始工作")
elif event == 'pause':
outer_state = 'paused'
print("暂停")
elif inner_state == 'working':
if event == 'complete':
inner_state = 'idle'
print("工作完成")
elif event == 'error':
inner_state = 'error'
print("出错")
elif inner_state == 'error':
if event == 'retry':
inner_state = 'working'
print("重试")
elif event == 'abort':
outer_state = 'terminated'
print("终止")
elif outer_state == 'paused':
if event == 'resume':
outer_state = 'active'
print("恢复")
elif outer_state == 'terminated':
if event == 'reset':
outer_state = 'active'
inner_state = 'idle'
print("重置")
双状态变量实现了内外两层状态机。外状态控制生命周期，内状态控制业务逻辑。这种模式在协议实现中很常见。
用yield from实现状态机子例程：
def idle_state():
while True:
event = yield
if event == 'start':
print("IDLE: 检测到start事件，切换到running")
return 'running'
print(f"IDLE: 忽略{event}事件")
def running_state():
count = 0
while True:
event = yield
if event == 'stop':
print("RUNNING: 检测到stop事件，切换到idle")
return 'idle'
elif event == 'tick':
count += 1
print(f"RUNNING: tick {count}")
def main_fsm():
state = 'idle'
idle_gen = idle_state()
running_gen = None
next(idle_gen)
while True:
event = yield
if state == 'idle':
result = idle_gen.send(event)
if result == 'running':
state = 'running'
running_gen = running_state()
next(running_gen)
elif state == 'running':
result = running_gen.send(event)
if result == 'idle':
state = 'idle'
idle_gen = idle_state()
next(idle_gen)
fsm = main_fsm()
next(fsm)
fsm.send('start')
fsm.send('tick')
fsm.send('tick')
fsm.send('stop')
每个状态有自己的生成器，状态切换时销毁旧生成器创建新生成器。这种方式隔离了不同状态的逻辑，避免了巨大的if-else链。
生成器状态机的性能特征也很重要。每个yield/send操作大约耗时0.1微秒，比函数调用慢一个数量级。对于每秒处理数万个事件的场景（如网络协议栈），这个开销可以接受。但对于高吞吐场景（每秒百万级事件），应该用基于表驱动的状态机。
Python 3.10引入的match语句让状态机实现更简洁：
class MatchFSM:
def __init__(self):
self.state = 'idle'
def handle(self, event):
match (self.state, event):
case ('idle', 'start'):
self.state = 'running'
print("开始")
case ('running', 'stop'):
self.state = 'idle'
print("停止")
case ('running', 'data' as e):
print(f"处理数据: {e}")
case (_, _):
print(f"忽略事件: {event}从状态{self.state}")
match语句的模式匹配把状态转换表的可读性提升了一个层次。每个case对应一个转换规则，_是通配符匹配任何值。