我可以修补 Python 的断言以获得 py.test 提供的输出吗?

Can I patch Python's assert to get the output that py.test provides?

Pytest 的失败断言输出比 Python 中的默认输出信息更丰富、更有用。我想在通常 运行 我的 Python 程序时利用它,而不仅仅是在执行测试时。有没有办法从我的脚本中覆盖 Python 的 assert 行为以使用 pytest 打印堆栈跟踪而不是同时 运行 我的程序仍然是 python script/pytest_assert.py

示例程序

def test_foo():
  foo = 12
  bar = 42
  assert foo == bar

if __name__ == '__main__':
  test_foo()

$ python script/pytest_assert.py

Traceback (most recent call last):
  File "script/pytest_assert.py", line 8, in <module>
    test_foo()
  File "script/pytest_assert.py", line 4, in test_foo
    assert foo == bar
AssertionError

$ pytest script/pytest_assert.py

======================== test session starts ========================
platform linux -- Python 3.5.3, pytest-3.3.2, py-1.5.2, pluggy-0.6.0
rootdir: /usr/local/google/home/danijar, inifile:
collected 1 item                                                    

script/pytest_assert.py F                                     [100%]

============================= FAILURES ==============================
_____________________________ test_foo ______________________________

    def test_foo():
      foo = 12
      bar = 42
>     assert foo == bar
E     assert 12 == 42

script/pytest_assert.py:4: AssertionError
===================== 1 failed in 0.02 seconds =====================

想要的结果

$ python script/pytest_assert.py

Traceback (most recent call last):
  File "script/pytest_assert.py", line 8, in <module>
    test_foo()

    def test_foo():
      foo = 12
      bar = 42
>     assert foo == bar
E     assert 12 == 42

script/pytest_assert.py:4: AssertionError

进度更新

我得到的最接近的是这个,但它只适用于那个函数中的断言并且垃圾邮件跟踪:

import ast
import inspect

from _pytest import assertion


def test_foo():
  foo = []
  foo.append(13)
  foo = foo[-1]
  bar = 42
  assert foo == bar, 'message'


if __name__ == '__main__':
  tree = ast.parse(inspect.getsource(test_foo))
  assertion.rewrite.rewrite_asserts(tree)
  code = compile(tree, '<name>', 'exec')
  ns = {}
  exec(code, ns)
  ns[test_foo.__name__]()

$ python script/pytest_assert.py
Traceback (most recent call last):
  File "script/pytest_assert.py", line 21, in <module>
    ns[test_foo.__name__]()
  File "<name>", line 6, in test_foo
AssertionError: message
assert 13 == 42

免责声明

虽然肯定有办法重用 pytest 代码以所需格式打印回溯,但您需要使用的东西不是 public API 的一部分,所以最终的解决方案将过于脆弱,需要调用不相关的 pytest 代码(用于初始化目的)并且可能会在包更新时中断。最好的办法是重写关键部分,以 pytest 代码为例。

注释

基本上,下面的概念验证代码做了三件事:

  1. 将默认的 sys.excepthook 替换为自定义的:这是更改默认回溯格式所必需的。示例:

    import sys
    
    orig_hook = sys.excepthook
    
    def myhook(*args):
        orig_hook(*args)
        print('hello world')
    
    if __name__ == '__main__':
        sys.excepthook = myhook
        raise ValueError()
    

    将输出:

    Traceback (most recent call last):
      File "example.py", line 11, in <module>
        raise ValueError()
    ValueError
    hello world
    
  2. 而不是hello world,将打印格式化的异常信息。为此,我们使用 ExceptionInfo.getrepr()

  3. 为了访问断言中的附加信息,pytest 重写了 assert 语句(您可以在 this old article). To achieve that, pytest registers a custom import hook as specified in PEP 302. The hook is the most problematic part as it is tightly coupled to Config object, also I noticed some module imports to cause problems (I guess it doesn't fail with pytest only because the modules are already imported when the hook is registered; will try to write a test that reproduces the issue on a pytest run and create a new issue). I would thus suggest to write a custom import hook that invokes the AssertionRewriter. 这个 AST tree walker class 是断言重写中必不可少的部分,而 AssertionRewritingHook 不是那么重要。

代码

so-51839452
├── hooks.py
├── main.py
└── pytest_assert.py

hooks.py

import sys

from pluggy import PluginManager
import _pytest.assertion.rewrite
from _pytest._code.code import ExceptionInfo
from _pytest.config import Config, PytestPluginManager


orig_excepthook = sys.excepthook

def _custom_excepthook(type, value, tb):
    orig_excepthook(type, value, tb)  # this is the original traceback printed
    # preparations for creation of pytest's exception info
    tb = tb.tb_next  # Skip *this* frame
    sys.last_type = type
    sys.last_value = value
    sys.last_traceback = tb

    info = ExceptionInfo(tup=(type, value, tb, ))

    # some of these params are configurable via pytest.ini
    # different params combination generates different output
    # e.g. style can be one of long|short|no|native
    params = {'funcargs': True, 'abspath': False, 'showlocals': False,
              'style': 'long', 'tbfilter': False, 'truncate_locals': True}
    print('------------------------------------')
    print(info.getrepr(**params))  # this is the exception info formatted
    del type, value, tb  # get rid of these in this frame


def _install_excepthook():
    sys.excepthook = _custom_excepthook


def _install_pytest_assertion_rewrite():
    # create minimal config stub so AssertionRewritingHook is happy
    pluginmanager = PytestPluginManager()
    config = Config(pluginmanager)
    config._parser._inidict['python_files'] = ('', '', [''])
    config._inicache = {'python_files': None, 'python_functions': None}
    config.inicfg = {}

    # these modules _have_ to be imported, or AssertionRewritingHook will complain
    import py._builtin
    import py._path.local
    import py._io.saferepr

    # call hook registration
    _pytest.assertion.install_importhook(config)

# convenience function
def install_hooks():
    _install_excepthook()
    _install_pytest_assertion_rewrite()

main.py

调用hooks.install_hooks()后,main.py将修改回溯打印。 install_hooks() 调用后导入的每个模块都将在导入时重写断言。

from hooks import install_hooks

install_hooks()

import pytest_assert


if __name__ == '__main__':
    pytest_assert.test_foo()

pytest_assert.py

def test_foo():
    foo = 12
    bar = 42
    assert foo == bar

示例输出

$ python main.py
Traceback (most recent call last):
  File "main.py", line 9, in <module>
    pytest_assert.test_foo()
  File "/Users/hoefling/projects/private/Whosebug/so-51839452/pytest_assert.py", line 4, in test_foo
    assert foo == bar
AssertionError
------------------------------------
def test_foo():
        foo = 12
        bar = 42
>       assert foo == bar
E       AssertionError

pytest_assert.py:4: AssertionError

总结

我会去写一个自己的 AssertionRewritingHook 版本,而不是所有不相关的 pytest 东西。然而 AssertionRewriter 看起来几乎可以重用;虽然它需要一个 Config 实例,但它仅用于警告打印,可以留给 None.

完成后,编写自己的函数来正确格式化异常,替换 sys.excepthook 就大功告成了。