基于Python statemachine模块的多级菜单代码优化咨询

核心优化思路 #

• 解耦逻辑：将菜单显示、输入处理、状态切换与业务动作分离，彻底消除嵌套循环与多层条件判断
• 配置化管理：用字典集中定义每个菜单的选项映射，新增/修改菜单只需调整配置，无需改动核心循环逻辑
• 统一主循环：依托状态机的当前状态动态处理交互，菜单层级扩展时代码结构始终保持清晰

重构后的代码 #

from statemachine import State, StateMachine

class MenuStateMachine(StateMachine):
    # 定义所有状态
    main_menu = State('Main Menu', initial=True)
    menu1 = State('Option1')
    menu2 = State('Option2')
    submenu1 = State('Sub-Menu1')

    # 定义状态转换规则
    select_menu1 = main_menu.to(menu1)
    select_menu2 = main_menu.to(menu2)
    select_submenu1 = menu1.to(submenu1)
    back_to_main_menu = menu1.to(main_menu) | menu2.to(main_menu)
    back_to_menu1 = submenu1.to(menu1)

    # 菜单配置：键为状态ID，值包含菜单标题、选项列表（选项标识、显示文本、对应动作/状态转换）
    menu_configs = {
        'main_menu': {
            'title': '=== Main Menu ===',
            'options': [
                ('1', 'Option1', select_menu1),
                ('2', 'Option2', select_menu2),
                ('3', 'Exit', exit)
            ]
        },
        'menu1': {
            'title': '=== Option1 Menu ===',
            'options': [
                ('1', 'SubOption1', select_submenu1),
                ('2', 'SubOption2', lambda: print('There is no menu yet, please choose another')),
                ('3', 'SubOption3', lambda: print('There is no menu yet, please choose another')),
                ('4', 'Previous Menu', back_to_main_menu)
            ]
        },
        'submenu1': {
            'title': '=== SubOption1 Menu ===',
            'options': [
                ('1', 'Act1', lambda: print('Do Something Here')),
                ('2', 'Act2', lambda: print('Do Something Here')),
                ('3', 'Act3', lambda: print('Do Something Here')),
                ('4', 'Previous Menu', back_to_menu1)
            ]
        },
        'menu2': {
            'title': '=== Option2 Menu ===',
            'options': [
                ('1', 'Previous Menu', back_to_main_menu)
            ]
        }
    }

    def show_menu(self):
        """根据当前状态渲染对应菜单"""
        config = self.menu_configs[self.current_state.id]
        print(config['title'])
        for key, text, _ in config['options']:
            print(f'[{key}] {text}')

    def handle_choice(self, choice):
        """处理用户输入，执行对应动作或状态转换"""
        config = self.menu_configs[self.current_state.id]
        for key, _, action in config['options']:
            if key == choice:
                # 区分状态转换动作和普通业务动作
                if hasattr(action, 'trigger'):
                    action.trigger(self)
                else:
                    action()
                return True
        print('Invalid choice. Please try again.')
        return False

# 初始化状态机实例
menu = MenuStateMachine()

# 统一主循环，处理所有交互
while True:
    menu.show_menu()
    choice = input('Enter Choice: ')
    menu.handle_choice(choice)

优化细节说明 #

1. 配置化菜单：所有菜单的标题、选项、对应动作都集中在menu_configs中，新增子菜单只需添加新的状态和配置项，无需修改循环逻辑
2. 统一交互入口：主循环仅负责显示当前菜单、获取输入、调用处理函数，彻底消除原代码的嵌套while和多层if-elif
3. 动作与状态分离：状态转换（如select_menu1）和业务动作（如打印提示）通过统一逻辑处理，代码结构更清晰
4. 高扩展性：后续新增任意层级的子菜单，只需定义新的State、状态转换规则，再在menu_configs中添加对应配置即可

内容的提问来源于stack exchange，提问作者ExHunter