如何在 python 脚本执行中禁用 Numpy 断言

Question

我正在 Bash 脚本内运行 Python 脚本。我在多个地方使用 NumPy 的

assert_allclose

和

assert

。我想通过禁用所有此类实例来运行此 Python 脚本，而不需要手动注释掉它们。

这是一个简单的例子：

python run.py

if __name__ == "__main__":
    assert False
    np.assert_allclose(1, 2)

我知道可以将

-O

标志添加到

python run.py

。但这似乎并没有禁用 NumPy 的断言，因为

np.assert_allclose

使用

raise AssertionError(msg)

而不是

assert

本身。

Answer 1

正如评论中提到的，我不认为你正在寻找的是内置的，但这里有一个你可以用来完成这项工作的包装器。

import sys

def wrap_assertion(f, enabled=True):
    def assertion(*args, **kwargs):
        if "disable_assertions" in sys.argv or not assertion.enabled:
            pass
        else:
            return f(*args, **kwargs)
    assertion.enabled = enabled
    return assertion

如果您在 IDE 中运行脚本，您可以控制是否启用断言，如下所示：

# Wrap np.testing.assert_allclose; disable by default
import numpy as np
assert_allclose = wrap_assertion(np.testing.assert_allclose, enabled=False)

# This does nothing
assert_allclose(1, 2)

# Enable the assertion
assert_allclose.enabled = True

# This raises an AssertionError
assert_allclose(2, 3)
# AssertionError:
# Not equal to tolerance rtol=1e-07, atol=0
# Mismatched elements: 1 / 1 (100%)
# Max absolute difference: 1
# Max relative difference: 0.33333333
#  x: array(2)
#  y: array(3)

或者，如果您从命令行运行 Python 脚本，则可以将

'disable_assertions'

作为命令行参数传递，以覆盖程序中指定的选择。例如，如果上述代码位于名为

run.py

的脚本中，则

python run.py

将导致

AssertionError

显示为注释，但

python run.py disable_assertions

将终止且不会出现错误。