如何在Python脚本运行时禁用NumPy断言及原生assert语句

问题 #

我在Bash里运行Python脚本，脚本里既有Python原生的assert语句，又用了NumPy的np.assert_allclose。想用一个方法一次性禁用所有这些断言，不用手动注释代码。试过python -O run.py，但这个参数只能禁用原生assert，管不了np.assert_allclose——因为它是主动抛出AssertionError，不是用原生assert语法写的。

解决办法 #

方法1：在脚本里加猴子补丁 #

在你的run.py最开头加上这段代码，它会在-O模式下把np.assert_allclose换成空函数，配合-O参数就能同时禁用两种断言：

import numpy as np
import sys

# 只有在-O优化模式下才替换断言函数
if sys.flags.optimize > 0:
    def no_op_assert(*args, **kwargs):
        pass
    np.assert_allclose = no_op_assert

然后运行命令：

python -O run.py

方法2：用独立补丁模块（不修改原脚本） #

如果不想改原脚本，先创建一个叫disable_numpy_asserts.py的文件，内容如下：

import numpy as np
import sys

if sys.flags.optimize > 0:
    # 把np.assert_allclose替换成啥也不做的lambda
    np.assert_allclose = lambda *args, **kwargs: None

然后用这个命令运行脚本：

python -O -m disable_numpy_asserts run.py

-m参数会先加载补丁模块，再运行你的脚本，这样原代码完全不用动。

额外提示 #

• sys.flags.optimize在-O模式下是1，-OO模式下是2，这个判断能确保补丁只在你需要禁用断言的时候生效，不影响平时调试。
• 如果脚本里还有其他NumPy断言（比如np.assert_equal），照着上面的方法替换就行。

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