如何在Jupyter Lab中实现类Pytest的详细断言输出？

下面是几种可行的实现方式，不需要修改笔记本反复运行就能获取断言失败的详细信息：

1. 使用better-exceptions库 #

该库会增强Python默认的异常输出，包括断言错误，能展示表达式各部分的具体取值，风格接近Pytest。

• 安装：

  pip install better-exceptions
  

• 在Jupyter Notebook/Lab中启用：

  import better_exceptions
  better_exceptions.hook()
  

  之后断言失败时，会自动显示更详细的上下文信息，比如断言表达式中各个变量的实际值。

2. 利用pytest的断言重写机制 #

直接借用pytest的断言增强能力，无需运行完整的pytest测试套件：

• 先确保已安装pytest：

  pip install pytest
  

• 在Notebook中执行以下代码启用断言重写：

  from _pytest.assertion.rewrite import rewrite_asserts
  # 对当前Notebook模块的断言进行重写
  rewrite_asserts(globals()['__file__'] if '__file__' in globals() else 'temp_notebook.py')
  

  重写后，你的普通assert语句会具备pytest风格的详细错误输出，比如展示断言两边的实际取值。

3. 使用assertpy库 #

这是一个专门优化断言体验的库，提供更直观的断言语法，同时输出清晰的错误信息：

• 安装：

  pip install assertpy
  

• 针对你的fold样本数验证场景，示例代码：

  from assertpy import assert_that
  
  # 假设fold_sizes是各fold的测试样本数量列表
  fold_sizes = [8, 12, 9]
  assert_that(fold_sizes).each_is_greater_than_or_equal_to(10)
  

  断言失败时会直接输出：

  AssertionError: Expected each element in <[8, 12, 9]> to be greater than or equal to <10>, but found <8> and <9>.
  

4. 自定义断言函数 #

如果不想依赖第三方库，可以自己编写简单的断言函数，在失败时输出针对性的详细信息：

def assert_fold_min_samples(fold_sizes, min_count=10):
    invalid = [idx for idx, size in enumerate(fold_sizes) if size < min_count]
    if invalid:
        raise AssertionError(f"第{', '.join(map(str, invalid))}个fold的测试样本数不足{min_count}个，对应数量为{[fold_sizes[i] for i in invalid]}")

# 使用示例
fold_sizes = [8, 12, 9]
assert_fold_min_samples(fold_sizes)

失败时会直接指出具体哪些fold不满足要求，无需重新修改代码排查问题。

内容的提问来源于stack exchange，提问作者David Davíd