如何实现Python脚本指定代码段的不可中断执行？

问题描述 #

需要让Python脚本中的指定代码块具备不可中断性：当触发键盘中断（Ctrl+C）时，程序会等待该代码块执行完成后，再响应中断。

具体场景：

• 执行高维KD空间划分操作，仅小部分区域为相关区域
• 程序定期将数据写入文件，后续根据文件内容决定下一步划分位置
• 程序支持中断重启，读取数据文件即可快速恢复，但写入阶段被中断会导致文件损坏，浪费大量计算时间
• 非代码层面要求，必须保证程序在其他区域可中断，仅指定代码块延迟响应中断

预想代码结构：

while continue_condition:
    process_data() # 此处触发中断会直接终止程序
    with dont_interrupt: 
        read_write_operations() # 此处触发中断会等待代码执行完毕再响应
    # 在dont_interrupt上下文内触发的中断会在此处抛出

核心需求：确保指定代码块完整执行后再处理中断，避免数据损坏是需求诱因，但不是核心目的。

解决方案：自定义上下文管理器 #

通过Python的contextlib.contextmanager实现dont_interrupt上下文管理器，捕获上下文内的KeyboardInterrupt异常，延迟到代码块执行完成后再抛出。

完整实现代码 #

from contextlib import contextmanager
import time

@contextmanager
def dont_interrupt():
    interrupt_triggered = False
    try:
        yield
    except KeyboardInterrupt:
        interrupt_triggered = True
    # 上下文退出后，若捕获到中断则重新抛出
    if interrupt_triggered:
        raise KeyboardInterrupt

# 模拟业务函数
def process_data():
    print("处理数据中（可中断）...")
    time.sleep(2)

def read_write_operations():
    print("执行文件读写操作（不可中断，延迟响应中断）...")
    time.sleep(3)
    print("读写操作完成")

# 主逻辑示例
if __name__ == "__main__":
    continue_condition = True
    try:
        while continue_condition:
            process_data()
            with dont_interrupt():
                read_write_operations()
            # 示例中仅执行一次循环，实际可根据业务逻辑设置终止条件
            continue_condition = False
    except KeyboardInterrupt:
        print("程序已安全响应中断并退出")

工作机制 #

1. 当在dont_interrupt上下文内按下Ctrl+C，KeyboardInterrupt异常会被捕获，标记interrupt_triggered为True
2. 上下文内的代码会继续执行直到完成，不会被中断
3. 上下文退出后，程序会重新抛出KeyboardInterrupt，此时程序在上下文外处理中断逻辑
4. 若在上下文外触发中断，程序会立即响应，不会延迟

注意要点 #

• 避免让上下文内的代码执行时间过长，否则用户需要等待较久才能终止程序
• 确保上下文内的操作具备完整性，即使中途触发中断，也能保证代码执行完后数据处于安全状态

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