我能否修补Python的assert来获得py.test提供的输出?

5

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

1
由于assert是一个关键字,答案将是“不行,除非重新编译解释器”。必须问的一个问题是:如果你喜欢pytest的输出,为什么不使用pytest呢? - DeepSpace
2
pytest除了处理“AssertionError”之外什么也不做。你应该尝试一下。 - Klaus D.
@KlausD。谢谢,这是个好主意。您知道如何调用pytest以打印这些细节吗?如果手动实现,我将不得不通过编程方式访问断言上面的源代码和变量值。 - danijar
@DeepSpace 我无法自由地决定在部署环境中运行我的脚本的可执行文件。由于pytest只是一个使用普通Python解释器执行的Python模块,因此必须能够在不重新编译的情况下完成。请参见上面的其他评论 :) - danijar
关于对这个问题过于宽泛的关闭投票,我很想听听背后的原因。我在问题中提供了一个明确的示例输出,说明了我想要实现的目标。 - danijar
1
你可能会在 https://github.com/pytest-dev/pytest/blob/master/src/_pytest/assertion/rewrite.py 中找到一些 pytest 的魔法。 - Klaus D.
1个回答

5

免责声明

虽然肯定有一种方法可以重用pytest代码以打印所需格式的回溯,但需要使用的内容不是公共API的一部分,因此生成的解决方案将过于脆弱,需要调用不相关的pytest代码(用于初始化),并且可能在软件包更新时出现故障。最好的方法是重写关键部分,使用pytest代码作为示例。

注释

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

  1. Replace the default sys.excepthook with the custom one: this is necessary to alter the default traceback formatting. Example:

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

    will output:

    Traceback (most recent call last):
      File "example.py", line 11, in <module>
        raise ValueError()
    ValueError
    hello world
    
  2. Instead of hello world, the formatted exception info will be printed. We use ExceptionInfo.getrepr() for that.

  3. To access the additional info in asserts, pytest rewrites the assert statements (you can get some rough info about how they look like after rewrite in 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. This AST tree walker class is the essential part in assertion rewriting, while the AssertionRewritingHook is not that important.

代码

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会修改回溯(traceback)的输出。在调用了install_hooks()函数之后导入的每个模块都将在导入时重写assert语句。

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/stackoverflow/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,就完成了。


哎呀,非常感谢!我不得不向配置存根中添加 config.getini = lambda _: [] 才能运行最新版本的 pytest。然而,我得到的输出并不完全正确,它显示了 > assert foo == bar\nE AssertionError 而不是 > assert foo=bar\nE assert 13 == 42,因此它没有解析变量值。 - danijar
我使用了一个不同的脚本并调用了 _pytest.assertion.rewrite.rewrite_asserts(tree) 在我想要运行的函数的AST上,最终获得了正确的输出。 我在我的问题中包含了这个示例。 但是,它不使用importhooks(我还没有完全理解)所以它只适用于该函数内部的断言。 - danijar
如果您不注册导入钩子,那么您将无法重写导入模块中的断言,因此重写仅适用于当前模块的字节码。该钩子将重写导入模块的字节码(请查看__pycache__目录,其中应包含带有重写断言的module_..._PYTEST.pyc文件),并使用此修改后的字节码代替解释器默认编译的内容。 - hoefling
是的,那很有道理。你知道如何编写自定义导入钩子来应用AssertionRewriter而不是调用pytest的install_importhook吗?我还不太确定为什么以E 开头的行在你的代码示例中打印出来的方式不同。 - danijar
导入钩子本质上相当简单;它是模块查找器和模块加载器的组合。以下是一些示例:onetwothreeAssertionRewritingHook 结合了两者,通过在一个类中实现 find_moduleload_module,这些内容包含了您需要在自己的实现中复制的最重要逻辑。 - hoefling

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接