Python脚本中灵活禁用NumPy及标准断言的策略与实践

# 示例：Python标准断言
if __name__ == "__main__":
    assert False, "This will be ignored with -O"

import numpy as np

# 示例：NumPy断言
if __name__ == "__main__":
    # np.assert_allclose(1, 2) # 这行在 -O 模式下依然会抛出 AssertionError
    pass

import sys
import numpy as np

def wrap_assertion(original_assertion_func, enabled_by_default=True):
    """
    创建一个断言包装器，允许在运行时控制断言的启用/禁用。

    Args:
        original_assertion_func: 原始的断言函数（例如 np.testing.assert_allclose）。
        enabled_by_default: 包装器创建时断言是否默认启用。

    Returns:
        一个新的断言函数，带有 .enabled 属性和命令行控制逻辑。
    """
    def assertion_wrapper(*args, **kwargs):
        # 检查包装器自身的 enabled 属性，以及命令行参数是否包含 'disable_assertions'
        # 如果包装器启用且命令行未指定禁用，则执行原始断言
        if assertion_wrapper.enabled and "disable_assertions" not in sys.argv:
            return original_assertion_func(*args, **kwargs)
        # 否则，断言被跳过，不执行任何操作
        return None

    # 为包装器函数添加一个可控制的 enabled 属性
    assertion_wrapper.enabled = enabled_by_default
    return assertion_wrapper

# 示例：包装 np.testing.assert_allclose
# 默认设置为禁用，除非通过代码或命令行显式启用
assert_allclose = wrap_assertion(np.testing.assert_allclose, enabled_by_default=False)

# 示例：包装一个自定义的断言函数
def my_custom_assert(condition, message="Custom assertion failed"):
    if not condition:
        raise AssertionError(message)

my_assert = wrap_assertion(my_custom_assert, enabled_by_default=True)

# run_script_internal.py
import numpy as np
import sys

# 假设 wrap_assertion 函数已定义如上
def wrap_assertion(original_assertion_func, enabled_by_default=True):
    def assertion_wrapper(*args, **kwargs):
        if assertion_wrapper.enabled and "disable_assertions" not in sys.argv:
            return original_assertion_func(*args, **kwargs)
        return None
    assertion_wrapper.enabled = enabled_by_default
    return assertion_wrapper

# 包装 np.testing.assert_allclose，默认禁用
assert_allclose = wrap_assertion(np.testing.assert_allclose, enabled_by_default=False)

print("--- 初始状态：断言默认禁用 ---")
try:
    assert_allclose(1, 2) # 这行代码不会抛出错误
    print("assert_allclose(1, 2) 未触发错误（预期）")
except AssertionError as e:
    print(f"assert_allclose(1, 2) 触发错误: {e} (非预期)")

# 显式启用断言
assert_allclose.enabled = True
print("\n--- 状态变更：断言已启用 ---")
try:
    assert_allclose(2, 3) # 这行代码会抛出 AssertionError
    print("assert_allclose(2, 3) 未触发错误 (非预期)")
except AssertionError as e:
    print(f"assert_allclose(2, 3) 触发错误: {e} (预期)")

# 再次禁用断言
assert_allclose.enabled = False
print("\n--- 状态变更：断言再次禁用 ---")
try:
    assert_allclose(4, 5) # 这行代码不会抛出错误
    print("assert_allclose(4, 5) 未触发错误（预期）")
except AssertionError as e:
    print(f"assert_allclose(4, 5) 触发错误: {e} (非预期)")

--- 初始状态：断言默认禁用 ---
assert_allclose(1, 2) 未触发错误（预期）

--- 状态变更：断言已启用 ---
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) (预期)

--- 状态变更：断言再次禁用 ---
assert_allclose(4, 5) 未触发错误（预期）

# run_script_cmd.py
import numpy as np
import sys

# 假设 wrap_assertion 函数已定义如上
def wrap_assertion(original_assertion_func, enabled_by_default=True):
    def assertion_wrapper(*args, **kwargs):
        if assertion_wrapper.enabled and "disable_assertions" not in sys.argv:
            return original_assertion_func(*args, **kwargs)
        return None
    assertion_wrapper.enabled = enabled_by_default
    return assertion_wrapper

# 包装 np.testing.assert_allclose，默认启用
assert_allclose = wrap_assertion(np.testing.assert_allclose, enabled_by_default=True)

# 包装一个标准Python断言，使其也受此机制控制
# 注意：对于Python内置的assert语句，更好的做法是使用-O标志，但这里演示包装器的通用性
def standard_assert(condition, msg="Assertion failed"):
    if not condition:
        raise AssertionError(msg)

my_assert = wrap_assertion(standard_assert, enabled_by_default=True)


print("--- 脚本开始执行 ---")

try:
    assert_allclose(1, 2)
    print("assert_allclose(1, 2) 未触发错误")
except AssertionError as e:
    print(f"assert_allclose(1, 2) 触发错误: {e}")

try:
    my_assert(False, "This is a wrapped standard assertion.")
    print("my_assert(False) 未触发错误")
except AssertionError as e:
    print(f"my_assert(False) 触发错误: {e}")

print("--- 脚本执行完毕 ---")