Assertion Rewriting

Overview

The **Assertion Rewriting** module provides a mechanism to automatically transform Python `assert` statements in test code into enhanced forms that produce detailed, user-friendly error messages when assertions fail. This rewriting occurs transparently during module import via a custom import hook that intercepts test module loading, modifies their abstract syntax trees (ASTs), and compiles the rewritten code for execution.

This feature addresses the common problem in testing where default assertion failures show only the expression that failed without detailed introspection into intermediate values or sub-expressions. By rewriting asserts, pytest can present rich failure explanations that help developers quickly understand why a test failed, improving debugging efficiency.

Core Concepts and Purpose

• Import Hook for Assertion Rewriting: The module implements a PEP 302/451-compliant import hook (AssertionRewritingHook) that intercepts the import of test modules matching configured filename patterns (e.g., test_*.py). It reads their source, rewrites all assert statements into expanded forms, and loads the transformed code instead of the original.

• AST Transformation for Assert Statements: The central transformation is performed by AssertionRewriter, which traverses the AST of each test module and replaces every assert statement with an equivalent block of code that:

  • Evaluates sub-expressions and intermediate values.

  • Collects these values to build an informative error message.

  • Raises an AssertionError with this detailed message if the assertion condition is false.

  • Optionally calls hooks if the assertion passes (experimental).

• Caching of Rewritten Bytecode: To optimize repeated test runs, rewritten modules are cached as .pyc files with a custom tag in pycache directories. This avoids repeated rewriting and speeds up imports while ensuring correctness.

• Selective Rewriting: The hook only rewrites test files or explicitly marked modules, based on filename patterns and command-line specified paths, avoiding unnecessary overhead on unrelated modules.

How Assertion Rewriting Works

1. Import Hook Interception

• When Python imports a module, pytest's AssertionRewritingHook participates in the import system by implementing find_spec(), create_module(), and exec_module() methods.

• It checks if the module should be rewritten according to configured test file patterns (python_files) and whether it is already rewritten.

• If rewriting is needed, it reads the source file and checks for a cached rewritten .pyc.

• If no valid cache exists, it calls _rewrite_test() to parse and transform the source.

2. AST Parsing and Transformation

• _rewrite_test() reads the source code and parses it into an AST using Python's ast module.

• It then invokes rewrite_asserts(), which creates an AssertionRewriter instance that walks the AST.

• AssertionRewriter.run() traverses the AST nodes and locates all assert statements.

• Each assert is replaced by a series of AST statements constructed by visit_Assert():

  • The original assert expression is recursively visited to capture sub-expressions, variables, and intermediate results.

  • New temporary variables are created to hold these intermediate values.

  • An if statement checks the negation of the original assert condition.

  • If false, an AssertionError is raised with a detailed message constructed from the collected sub-expressions.

  • Temporary variables are cleared after use to avoid namespace pollution.

3. Enhanced Assertion Messages

• The rewritten assertions use helper functions (e.g., _saferepr, _format_explanation) to safely obtain string representations of values, escape special characters, and format the failure message.

• The detailed message shows the original expression, sub-expression results, and their values, providing a clear explanation of what failed.

4. Execution of Rewritten Code

• The rewritten AST is compiled to a code object.

• This code object is saved to a .pyc cache file with a pytest-specific tag to avoid conflicts.

• The rewritten module is executed in place of the original, enabling enhanced assertion introspection during test runs without requiring changes to test source code.

Interaction with Other Parts of the System

• Configuration (Config): The import hook uses pytest configuration to get filename patterns for test files and options controlling assertion rewriting behavior (e.g., enabling assertion pass hooks).

• Session: The hook can be associated with a test session to optimize rewriting decisions based on test paths specified on the command line.

• Assertion Utilities (_pytest.assertion.util): The rewritten code leverages utilities for formatting assertion failure explanations and managing assertion pass hooks.

• Safe Representation (_pytest._io.saferepr): Used to generate safe, truncated, and newline-escaped string representations of objects for inclusion in failure messages.

• Warning System: The hook warns if a module was imported before rewriting, which can prevent rewriting and degrade failure message quality.

Key Functional Components

AssertionRewritingHook (Import Hook)

• Implements the import system interface to selectively rewrite source modules during import.

• Maintains caches and flags to avoid infinite recursion or redundant rewriting.

• Uses filename patterns and session paths to decide whether to rewrite a module.

• Reads source, rewrites asserts, caches compiled bytecode, and executes rewritten code.

AssertionRewriter (AST Transformer)

• Visits all AST nodes in a module, focusing on assert statements.

• For each assert:

  • Visits sub-expressions recursively (visit_Name, visit_BoolOp, visit_BinOp, etc.).

  • Generates temporary variables to hold intermediate values.

  • Builds a detailed explanation string with placeholders and formatting parameters.

  • Replaces the assert with an if statement that raises an AssertionError with the enriched message if the assert fails.

• Handles specific AST constructs like the walrus operator (NamedExpr), comparisons, calls, attribute access, boolean operations, unary and binary operations, and others to extract detailed runtime information.

Helper Functions and Utilities

• _saferepr: Safely represent objects with escaped newlines and size limits.

• _format_assertmsg: Format custom assertion messages preserving newlines.

• _get_assertion_exprs: Tokenizes source to map line numbers to the original assert expression strings.

• _write_pyc, _read_pyc: Manage caching of rewritten bytecode files.

Example: How an assert Statement is Rewritten

Original code:

assert x + y == z

Rewritten code (conceptualized):

# Evaluate sub-expressions and assign to temporary variables
@py_assert0 = x
@py_assert1 = y
@py_assert2 = @py_assert0 + @py_assert1
@py_assert3 = z

# Build explanation string with intermediate values
explanation = "(@py_assert0 + @py_assert1) == @py_assert3\n({0} + {1} == {2})".format(
    repr(@py_assert0), repr(@py_assert1), repr(@py_assert3)
)

# Check condition and raise detailed AssertionError if false
if not (@py_assert2 == @py_assert3):
    raise AssertionError("assert " + explanation)

# Clear temporary variables
@py_assert0 = None
@py_assert1 = None
@py_assert2 = None
@py_assert3 = None

This expanded form provides explicit feedback on the values of `x`, `y`, `z`, and the result of `x + y`, which aids debugging.

Visual Diagram: Assertion Rewriting Workflow

sequenceDiagram
    participant Importer as Python Import System
    participant Hook as AssertionRewritingHook
    participant AST as AssertionRewriter (AST Transformer)
    participant Cache as Bytecode Cache (.pyc)
    participant TestModule as Test Module Execution

    Importer->>Hook: Request to import module
    Hook->>Hook: Decide if rewriting needed
    alt Rewrite required
        Hook->>Cache: Check for cached rewritten bytecode
        alt Cache hit
            Hook->>TestModule: Load and execute cached bytecode
        else Cache miss
            Hook->>Hook: Read source file
            Hook->>AST: Parse and rewrite asserts in AST
            AST->>Hook: Return rewritten AST
            Hook->>Hook: Compile rewritten AST to bytecode
            Hook->>Cache: Write rewritten bytecode to cache
            Hook->>TestModule: Execute rewritten bytecode
        end
    else No rewrite
        Importer->>TestModule: Load and execute original bytecode
    end

Summary of Assertion Rewriting Module

• Provides an import hook that transparently rewrites assert statements in test modules.

• Transforms asserts into expanded code that captures intermediate values and generates detailed failure messages.

• Enhances test debugging by providing rich introspection without requiring test code changes.

• Uses AST manipulation, caching strategies, and tightly integrates with pytest's configuration and assertion utility functions.

• Balances performance and utility by selectively rewriting only test files and caching rewritten bytecode on disk.