如何在调用multiprocessing.Pool前检测用户是否添加__name__ == '__main__'保护

可行性结论 #

完全可行，无需修改Python解释器逻辑，仅通过包内的前置判断即可实现。

实现思路 #

• Windows系统下multiprocessing默认使用spawn方式启动子进程，子进程启动时会全量导入入口脚本
• 若用户未加if __name__ == '__main__'保护，子进程导入入口脚本时会重新执行my_function调用逻辑
• 我们只需要在my_function启动Pool前，判断当前进程是否为子进程，如果是，说明用户未加保护，直接抛出明确错误即可

代码实现 #

修改后的my_module.py参考如下：

# my_module.py
import sys
from multiprocessing import Pool, current_process

def f(x):
    return x*x

def my_function(x_max):
    # 仅在spawn启动模式下需要检测，fork/forkserver模式无此问题
    if sys.platform == 'win32' or current_process()._config.get('start_method') == 'spawn':
        # 主进程name固定为MainProcess，子进程name格式为SpawnPoolWorker-x
        if current_process().name != 'MainProcess':
            raise RuntimeError(
                "检测到入口脚本未添加`if __name__ == '__main__'`保护语句，\n"
                "这会导致multiprocessing子进程启动时递归执行代码，进程挂起。\n"
                "请调整您的代码，将my_function的调用放在`if __name__ == '__main__'`代码块内。"
            )
    with Pool(5) as p:
        print(p.map(f, range(x_max)))

效果验证 #

粗心用户运行未加保护的脚本时，不会再无限制挂起，会立刻收到上述错误提示，可清晰定位修改点。

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