动态禁用Python脚本中NumPy断言的实用方法

import sys

def wrap_assertion(f, enabled=True):
    """
    创建一个断言函数的包装器，允许动态启用或禁用该断言。

    Args:
        f (callable): 原始的断言函数（例如 np.testing.assert_allclose）。
        enabled (bool): 初始状态，True 表示启用，False 表示禁用。

    Returns:
        callable: 被包装后的断言函数。
    """
    def assertion(*args, **kwargs):
        # 检查包装器的启用状态以及命令行参数
        # 如果包装器启用且命令行参数中不包含 'disable_assertions'，则执行原始断言
        if assertion.enabled and "disable_assertions" not in sys.argv:
            return f(*args, **kwargs)
    assertion.enabled = enabled # 将启用状态作为属性附加到包装器函数上
    return assertion

import numpy as np
# 假设 wrap_assertion 函数已定义在同一个文件或可导入模块中

# 1. 包装 np.testing.assert_allclose，默认禁用
# 注意：这里需要从 np.testing 导入原始函数
my_assert_allclose = wrap_assertion(np.testing.assert_allclose, enabled=False)

print("--- 默认禁用状态 ---")
try:
    my_assert_allclose(1, 2) # 此时不会引发 AssertionError
    print("my_assert_allclose(1, 2) 已被禁用，未引发错误。")
except AssertionError as e:
    print(f"发生错误 (不应出现): {e}")

# 2. 启用断言
my_assert_allclose.enabled = True
print("\n--- 启用状态 ---")
try:
    my_assert_allclose(2, 3) # 此时会引发 AssertionError
    print("my_assert_allclose(2, 3) 已启用，未引发错误 (不应出现)。")
except AssertionError as e:
    print(f"my_assert_allclose(2, 3) 已启用，成功捕获错误: {e}")

# 3. 再次禁用断言
my_assert_allclose.enabled = False
print("\n--- 再次禁用状态 ---")
try:
    my_assert_allclose(4, 5) # 此时再次被禁用
    print("my_assert_allclose(4, 5) 已再次禁用，未引发错误。")
except AssertionError as e:
    print(f"发生错误 (不应出现): {e}")

--- 默认禁用状态 ---
my_assert_allclose(1, 2) 已被禁用，未引发错误。

--- 启用状态 ---
my_assert_allclose(2, 3) 已启用，成功捕获错误:
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)

--- 再次禁用状态 ---
my_assert_allclose(4, 5) 已再次禁用，未引发错误。

# run_with_assertions.py
import sys
import numpy as np

# 假设 wrap_assertion 函数已定义在同一个文件或可导入模块中
def wrap_assertion(f, enabled=True):
    def assertion(*args, **kwargs):
        if assertion.enabled and "disable_assertions" not in sys.argv:
            return f(*args, **kwargs)
    assertion.enabled = enabled
    return assertion

# 包装 NumPy 的 assert_allclose，默认启用
my_assert_allclose = wrap_assertion(np.testing.assert_allclose, enabled=True)

# 包装一个模拟的自定义断言函数，用于演示通用性
def custom_assert(condition, message="Assertion failed"):
    if not condition:
        raise AssertionError(message)
my_assert = wrap_assertion(custom_assert, enabled=True)


if __name__ == "__main__":
    print("--- 脚本开始执行 ---")
    try:
        # 这是一个默认会失败的 NumPy 断言
        my_assert_allclose(1, 2)
        print("NumPy assert_allclose(1, 2) 未引发错误。")
    except AssertionError as e:
        print(f"NumPy assert_allclose(1, 2) 引发错误: {e}")

    try:
        # 这是一个默认会失败的自定义断言
        my_assert(False, "This custom assert should fail.")
        print("Custom assert(False) 未引发错误。")
    except AssertionError as e:
        print(f"Custom assert(False) 引发错误: {e}")

    print("--- 脚本执行结束 ---")