如何重载/重写Python的assert语句以实现断言通过时打印信息

嘿，完全懂你的需求——不想搞个自定义的my_assert函数替代原生语句，就想让原来的assert既能在失败时抛出带消息的AssertionError，又能在成功时输出提示信息对吧？下面给你两个实用的方案，都能满足你的要求：

方案一：上下文管理器（简单可靠，推荐） #

这个方法用上下文管理器包裹assert语句，既保留原生assert的语法和错误行为，又能轻松添加成功提示。

实现代码 #

from contextlib import contextmanager

@contextmanager
def assert_with_success(pass_msg):
    try:
        # 执行上下文内的代码（也就是你的assert语句）
        yield
    except AssertionError as e:
        # 断言失败时，先打印失败提示，再原样抛出异常，保留原生Traceback结构
        print(f"❌ Assertion failed: {e}")
        raise
    else:
        # 断言成功时，打印自定义的成功消息
        print(f"✅ Assertion passed: {pass_msg}")

使用示例 #

# 验证加法运算
with assert_with_success("1+1的结果符合预期"):
    assert 1 + 1 == 2, "加法运算错误"

# 验证列表非空
with assert_with_success("目标列表不为空"):
    assert len([1, 2, 3]) > 0, "列表为空"

# 故意触发失败，看原生结构是否保留
with assert_with_success("这个断言会失败"):
    assert 0, "测试失败消息"

优点 #

• 完全保留原生assert的语法和错误抛出逻辑，失败时的Traceback和原生一模一样
• 代码直观，不需要复杂的全局修改，只需要给需要提示的assert加个上下文包裹就行
• 成功消息和失败消息分离，逻辑清晰

方案二：全局跟踪（完全原生语法，适合调试） #

如果你连上下文包裹都不想加，就想完全用assert cond, msg的原生写法，那可以用代码跟踪的方式，解析你传入的消息元组来实现成功提示。

实现代码 #

import sys
import inspect

def _trace_assert(frame, event, arg):
    if event == 'line':
        # 获取当前执行的代码行
        code = frame.f_code
        line_idx = frame.f_lineno - 1  # 代码行是从0开始索引的
        try:
            lines = inspect.getsource(code).splitlines()
            if line_idx < len(lines):
                line_content = lines[line_idx].strip()
                # 判断当前行是否是assert语句
                if line_content.startswith('assert'):
                    # 拆分断言条件和消息部分
                    parts = line_content.split(',', 1)
                    if len(parts) == 2:
                        msg_segment = parts[1].strip()
                        # 假设用户传入的是(失败消息, 成功消息)的元组
                        if msg_segment.startswith('(') and msg_segment.endswith(')'):
                            try:
                                fail_msg, pass_msg = eval(msg_segment, frame.f_globals, frame.f_locals)
                                # 把成功消息存在当前帧的局部变量里
                                frame.f_locals['__pass_msg__'] = pass_msg
                                # 临时替换当前帧的断言错误处理，成功时打印消息
                                original_trace = frame.f_trace
                                def _check_success(frame_inner, event_inner, arg_inner):
                                    if event_inner == 'return':
                                        if '__pass_msg__' in frame_inner.f_locals:
                                            print(f"✅ Assertion passed: {frame_inner.f_locals['__pass_msg__']}")
                                    return original_trace
                                frame.f_trace = _check_success
                            except:
                                # 解析失败就忽略，保持原生行为
                                pass
        except:
            # 获取源码失败（比如交互模式）就跳过
            pass
    return _trace_assert

def enable_verbose_assert():
    """启用带成功提示的断言功能"""
    sys.settrace(_trace_assert)

使用示例 #

# 先启用功能
enable_verbose_assert()

# 直接用原生assert语法，消息传(失败消息, 成功消息)的元组
assert 1 + 1 == 2, ("加法运算错误", "1+1结果正确")
# 成功时会打印：✅ Assertion passed: 1+1结果正确

assert len([]) > 0, ("列表为空", "列表非空")
# 失败时抛出原生的AssertionError: 列表为空，Traceback结构不变

注意事项 #

• 这个方案依赖代码跟踪，会有轻微的性能损耗，适合调试阶段使用，上线前可以关闭
• 需要你传入元组作为消息，否则不会触发成功提示
• 在Python交互环境中可能无法获取源码，导致部分功能失效

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