Can I patch Python's assert to get the output

Question

Pytest's output for failed asserts is much more informative and useful than the default in Python. I would like to leverage this when normally running my Python program, not just when executing tests. Is there a way to, from within my script, overwrite Python's assert behavior to use pytest to print the stacktrace instead while still running my program as python script/pytest_assert.py?

Example program

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 =====================

Desired result

$ 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

Progress update

The closest I've got is this but it only works for asserts within that one function and spams the trace:

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

Answer 1

Disclaimer

Although there is surely a way to reuse pytest code to print the traceback in the desired format, stuff you need to use is not part of public API, so the resulting solution will be too fragile, require invocation of non-related pytest code (for initialization purposes) and likely break on package updates. Best bet would be rewriting crucial parts, using pytest code as an example.

Notes

Basically, the proof-of-concept code below does three things:

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.

Code

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

After calling hooks.install_hooks(), main.py will have modified traceback printing. Every module imported after install_hooks() call will have asserts rewritten on import.

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

Example output

$ 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

Summarizing

I would go with writing an own version of AssertionRewritingHook, without the whole non-related pytest stuff. The AssertionRewriter however looks pretty much reusable; although it requires a Config instance, it is only used for warning printing and can be left to None.

Once you have that, write your own function that formats the exception properly, replace sys.excepthook and you're done.