Behind the scenes of py.test's new assertion rewriting

py.test 2.1 was just released. py.test, which uses the Python assert statement to check test conditions, has long had support for displaying intermediate values in subexpressions of a failing assert statement. This feature is called assertion introspection. Historically, py.test performed assertion introspection by reinterpreting failed assertions in order to glean information about subexpressions. In assertion reinterpreting, py.test actually reruns the assertion noting intermediate values during interpretation. This works pretty well but is subject to several problems, most importantly that assert statements with side-effects can produce strange results because they are evaluated twice on failure. py.test 2.1's main new feature, which I wrote (with generous sponsorship from Merlinux GmbH), is a new assertion introspection technique called assertion rewriting. Assertion rewriting modifies the AST of test modules to produce subexpression information when assertions fail. This blog post will give a peek into how this is done and what the rewritten tests look like.

py.test tries to rewrite every module that it collects as a test module. Assertion rewriting uses a PEP 302 import hook to capture test modules for rewriting. I'm happy to report doing this was easier than I expected. Most of the code in the import hook I had to write was dealing with detecting test modules rather than supporting import's extremely complicated API. Rewriting has a non-zero cost during test collection, so py.test compiles rewritten modules to bytecode and caches them in the PEP 3147 PYC repository, __pycache__. One major thing I did have to account for was the possibility that multiple py.test processes would be writing PYC files. (This is a very real possibility when the xdist plugin is being used. Therefore, py.test uses only atomic operations on the rewritten PYC file. Windows, lacking atomic rename, was a pain here.

I'm now going to demonstrate what py.test's rewriting phase does to a test module. Let's dive in with a failing test for a (broken) function that is supposed to create empty files:

import os

def make_empty_file(name):
    with open(name, "w") as fp:
        fp.write("hello")

def test_make_empty_file():
    name = "/tmp/empty_test"
    make_empty_file(name)
    with open(name, "r") as fp:
        assert not fp.read()

This test nicely demonstrates the problem with py.test's old assertion method mentioned in the first paragraph. If we force the old assertion interpretation mode with --assert=reinterp, we see:

def test_make_empty_file():
name = "/tmp/empty_test"
make_empty_file(name)
with open(name, "r") as fp:
> assert not fp.read()
E AssertionError: (assertion failed, but when it was re-run for printing intermediate values, it did not fail. Suggestions: compute assert expression before the assert or use --no-assert)

test_empty_file.py:11: AssertionError

The problem is that assert statement has the side-effect of reading the file. When py.test reinterprets the assert statement, it uses the same file object, now at EOF, and read() returns an empty string. py.test's new rewriting mode fixes this by scanning the assert for introspection information before executing the test. Running py.test with assertion rewriting enabled gives the desired result:

def test_make_empty_file():
name = "/tmp/empty_test"
make_empty_file(name)
with open(name, "r") as fp:
> assert not fp.read()
E assert not 'hello'
E + where 'hello' = ()
E + where = .read

test_empty_file.py:11: AssertionError

So what magic has py.test worked to display such nice debugging information? This is what Python is actually executing:

def test_make_empty_file():
name = '/tmp/empty_test'
make_empty_file(name)
with open(name, 'r') as fp:
@py_assert1 = fp.read
@py_assert3 = @py_assert1()
@py_assert5 = (not @py_assert3)
if (not @py_assert5):
@py_format6 = ('assert not %(py4)s\n{%(py4)s = %(py2)s\n{%(py2)s = %(py0)s.read\n}()\n}' %
{'py0': (@pytest_ar._saferepr(fp) if ('fp' in @py_builtins.locals() is not @py_builtins.globals()) else 'fp'),
'py2': @pytest_ar._saferepr(@py_assert1),
'py4': @pytest_ar._saferepr(@py_assert3)})
raise AssertionError(@pytest_ar._format_explanation(@py_format6))
del @py_assert5, @py_assert1, @py_assert3

As you can see, it's not going to be winning any awards for beautiful Python! (Ideally, though, you'll never have to see or think about it.) Examining the rewritten code, we see a lot of internal variables starting with "@" have been created. The "@", invalid in Python identifiers, is to make sure internal names don't conflict with any user-defined names which might be in the scope. In the first four written lines under the with statement, the test of the assert statement has been expanded into its component subexpressions. This allows py.test to display the values of subexpressions should the assertion fail. If the assertion fails, the if statement in the fifth line of rewriting evaluates to True and a AssertionError will be raised. Under the if statement is the real mess. This is where the helpful error message is generated. The line starting with @py_format6 is simply does string formatting (with %) on a template generated from the structure of the assert statement. This template is filled in with the intermediate values of the expressions collected above. @py_builtins is the builtins module, used in case the test is shadowing builtins the rewriting code uses. The @pytest_ar variable is a special module of assertion formatting helpers. For example, @pytest_ar._saferepr is like builtin repr but gracefully handles long reprs and __repr__ methods that raise exceptions. A non-obvious trick in the format dict is the expression @pytest_ar._saferepr(fp) if ('fp' in @py_builtins.locals() is not @py_builtins.globals()) else 'fp'. This checks whether fp is a local variable or not and customizes the display accordingly. After the initial formatting, the helper function _format_explanation is called. This function produces the indentation and "+" you see in the error message. Finally, we note that if the assertion doesn't fail, py.test cleans up after itself by deleting temporary variables.

The example above is a fairly tame (and luckily also typical) assertion. Rewriting gets more "exciting" when boolean operations and comparisons enter because they require short circuit evaluation, which complicates both the expansion of expressions and formatting (think lots of nested ifs).

In conclusion, py.test's new assertion rewriting fixes some long standing issues with assertion introspection and continues py.test's long tradition of excellent debugging support. (There are now three(!) assertion introspection methods in py.test: two reinterpretation implementations as well as rewriting) I just hope I haven't scared you completely off py.test! :)