rewrite.py

Overview

The [rewrite.py](/projects/286/67427) module is a core component of the **pytest** testing framework responsible for **assertion rewriting**. Its primary goal is to transform Python `assert` statements in test modules into enhanced forms that provide rich, detailed error messages upon assertion failures. This transformation is achieved by hooking into Python’s import system, parsing test source code into an Abstract Syntax Tree (AST), rewriting `assert` statements to collect intermediate sub-expression values, and compiling the rewritten code for execution.

This mechanism allows pytest to display informative, introspective assertion failure explanations without requiring any changes from the test author. It significantly improves the debugging experience by showing the values of sub-expressions involved in failed assertions.

Key Components

1. AssertionRewritingHook (Import Hook)

• Implements Python's import system interfaces (MetaPathFinder and Loader) to intercept imports of test modules.

• Decides whether a module should be rewritten based on filename patterns (e.g., test_*.py), command-line specified test paths, or explicit markings.

• Reads the source code of modules to be rewritten.

• Checks and uses a cache of rewritten bytecode (.pyc) files to avoid unnecessary rewriting.

• If no valid cache exists, calls the rewriting process to transform the AST of the module.

• Compiles the rewritten AST back into bytecode and executes it.

• Manages caching and concurrency concerns when writing .pyc files.

• Provides methods to mark modules explicitly for rewriting and warns if a module was imported before rewriting, which can prevent rewriting.

2. AssertionRewriter (AST Transformer)

• Performs detailed AST transformations on assert statements found in test modules.

• Recursively visits AST nodes inside each assert expression to:

  • Introduce temporary variables to store intermediate values.

  • Build a detailed explanation string that includes the original expression and sub-expression values.

  • Replace the original assert with an if statement that raises a detailed AssertionError if the assertion fails.

• Supports advanced Python constructs like the walrus operator (NamedExpr), boolean operations, comparisons, attribute access, function calls, and unary/binary operations.

• Optionally supports the pytest_assertion_pass hook to invoke custom logic when assertions pass.

• Inserts import statements for helper functions at the top of rewritten modules.

• Ensures temporary variables are cleared after the assertion to avoid polluting namespace.

3. Helper Functions

• _rewrite_test: Parses source code into AST, invokes assertion rewriting, and compiles the rewritten AST.

• _read_pyc / _write_pyc: Manage reading and writing of cached rewritten bytecode files.

• _saferepr, _format_assertmsg, and other formatting helpers: Provide safe and detailed string representations for assertion messages.

• try_makedirs and get_cache_dir: Utilities for managing the .pyc cache directory.

• _get_assertion_exprs: Parses source code tokens to extract original assertion expressions by line number for use in messages.

Detailed Explanation of Classes and Functions

Class: AssertionRewritingHook

**Purpose:** Acts as a PEP 302/451 import hook that intercepts module imports and rewrites `assert` statements in test modules to enhance error reporting.

**Key Methods:**

• init(self, config: Config) -> None
  Initializes the hook with pytest configuration, loading filename patterns for test files.

• set_session(self, session: Session | None) -> None
  Associates a pytest test session to optimize path checks.

• find_spec(self, name: str, path: Sequence[str | bytes] | None = None, target: types.ModuleType | None = None) -> importlib.machinery.ModuleSpec | None
  Determines if a module should be rewritten. Returns a new module spec with this loader if rewriting is needed, otherwise None.

• create_module(self, spec: importlib.machinery.ModuleSpec) -> types.ModuleType | None
  Uses default module creation behavior (returns None).

• exec_module(self, module: types.ModuleType) -> None
  Performs the actual rewriting:

  • Reads cached rewritten bytecode if available and valid.

  • Otherwise, rewrites the source AST, compiles it, caches the bytecode, and executes it in the module namespace.

• mark_rewrite(self, *names: str) -> None
  Marks modules by name to be forcibly rewritten on import, including nested modules.

**Usage Example:**

hook = AssertionRewritingHook(config)
hook.mark_rewrite("my_test_module")
spec = hook.find_spec("my_test_module")
if spec:
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)

Class: AssertionRewriter (inherits ast.NodeVisitor)

**Purpose:** Transforms `assert` statements in an AST to enhanced versions that provide detailed failure explanations with intermediate values.

**Constructor:**

def __init__(self, module_path: str | None, config: Config | None, source: bytes) -> None:

• module_path: Path of the module being rewritten (used for warnings).

• config: pytest configuration for options like assertion pass hook.

• source: Source bytes of the module.

**Main Method:**

def run(self, mod: ast.Module) -> None:

• Entry point that traverses the entire AST and rewrites all assert statements.

• Inserts necessary helper imports at the top of the module.

• Manages scope tracking to properly handle nested functions/classes.

**Key Method:**

def visit_Assert(self, assert_: ast.Assert) -> list[ast.stmt]:

• Rewrites a single assert statement.

• Creates temporary variables for intermediate values.

• Builds a detailed explanation string using placeholders and formatting.

• Generates an if statement that raises a detailed AssertionError if the assertion fails.

• Supports optional pytest_assertion_pass hook calls.

• Clears temporary variables after use.

**Other Notable Visit Methods:**

• visit_Name, visit_NamedExpr: Handle variable names and walrus operator targets, applying special repr formatting.

• visit_BoolOp, visit_BinOp, visit_UnaryOp: Handle boolean, binary, and unary operations, capturing and formatting sub-expressions.

• visit_Compare: Handles chained comparisons, creating explanations for each operation.

• visit_Call, visit_Attribute, visit_Starred: Handle function calls, attribute access, and starred expressions with detailed formatting.

**Helper Methods:**

• variable() -> str: Generates a new unique temporary variable name.

• assign(expr: ast.expr) -> ast.Name: Assigns an expression to a new temporary variable and returns a reference to it.

• explanation_param(expr: ast.expr) -> str: Creates a named placeholder for string formatting explanations.

• push_format_context() / pop_format_context(expl_expr: ast.expr) -> ast.Name: Manage nested formatting contexts to build assertion messages.

**Usage Example:**

rewriter = AssertionRewriter(module_path="test_sample.py", config=pytest_config, source=source_bytes)
rewriter.run(ast_tree)
# The ast_tree now contains rewritten assert statements

Function: rewrite_asserts

def rewrite_asserts(
    mod: ast.Module,
    source: bytes,
    module_path: str | None = None,
    config: Config | None = None,
) -> None:

• Public API to rewrite all assert statements in a module AST.

• Instantiates AssertionRewriter and runs the rewriting process.

• Used internally by _rewrite_test.

Function: _rewrite_test

def _rewrite_test(fn: Path, config: Config) -> tuple[os.stat_result, types.CodeType]:

• Reads the source file at fn.

• Parses it into an AST.

• Calls rewrite_asserts to transform assert statements.

• Compiles the rewritten AST into a code object.

• Returns the file's stat info and the compiled code object.

Function: _read_pyc

def _read_pyc(
    source: Path, pyc: Path, trace: Callable[[str], None] = lambda x: None
) -> types.CodeType | None:

• Attempts to read a cached rewritten .pyc file.

• Validates the .pyc header against the source file's modification time and size.

• Returns the code object if valid, else None.

Function: _write_pyc

def _write_pyc(
    state: AssertionState,
    co: types.CodeType,
    source_stat: os.stat_result,
    pyc: Path,
) -> bool:

• Writes a compiled code object co to a .pyc file at pyc path.

• Uses atomic file replacement to avoid partial writes.

• Returns True if successful, else False.

Utility Functions

• _saferepr(obj: object) -> str: Produces a safe string representation suitable for assertion messages.

• _format_assertmsg(obj: object) -> str: Formats custom assertion messages preserving newlines.

• _get_assertion_exprs(src: bytes) -> dict[int, str]: Tokenizes source code to map line numbers to assertion test expressions.

• try_makedirs(cache_dir: Path) -> bool: Safely creates cache directories, handling errors like read-only filesystems.

• get_cache_dir(file_path: Path) -> Path: Returns the appropriate .pyc cache directory for a given source file.

Important Implementation Details and Algorithms

• Import Hook Logic:
  The AssertionRewritingHook integrates with Python's import machinery by implementing find_spec, create_module, and exec_module. It uses filename patterns and session test paths to decide whether to rewrite a module. It avoids rewriting .pyc files and modules already rewritten or marked to skip rewriting.

• AST Transformation:
  The AssertionRewriter traverses the AST in a depth-first manner. When it encounters an assert statement, it rewrites it by:

  • Creating temporary variables for sub-expressions to capture their values.

  • Constructing detailed explanation strings with placeholders filled by these variables.

  • Replacing the assert with an equivalent if not condition: raise AssertionError(...) block.

  • Managing variable scope and clearing temporary variables after assertion to avoid polluting the namespace.

• Caching Strategy:
  Rewritten bytecode is cached in .pyc files under __pycache__ or a custom cache prefix directory, with filenames tagged by Python implementation version and pytest version. This avoids repeated rewriting and speeds up imports.

• Safe Representation:
  The module employs specialized safe repr functions that escape newlines and limit length to ensure assertion messages are readable and do not break formatting.

• Support for New Python Features:
  The code handles Python 3.8+ features like the walrus operator (NamedExpr) and adapts imports for different Python versions (e.g., resource readers).

• Assertion Pass Hook (Experimental):
  When enabled, the rewritten assertions also call a pytest hook on passing assertions, allowing plugins to react to successful assertions.

• Concurrency and Atomicity:
  Writing .pyc files uses temporary files and atomic rename operations to avoid race conditions during parallel test runs.

Interaction with Other Parts of the System

• pytest Configuration (Config): Controls patterns for test file detection, verbosity for saferepr size, and enabling of assertion pass hooks.

• pytest Session (Session): Provides test paths specified on the command-line for optimized rewriting decisions.

• Assertion Utilities (_pytest.assertion.util): Provides formatting functions and hook implementations used by rewritten assertions.

• Safe Representation (_pytest._io.saferepr): Used to produce safe and concise string representations of objects for error messages.

• Warning System: Issues warnings if modules are imported before rewriting, which can degrade error message quality.

• Bytecode Import System: Integrates deeply with Python's importlib machinery to transparently rewrite and load test modules.

Visual Diagram: Class Structure and Main Workflow

classDiagram
    class AssertionRewritingHook {
        -config: Config
        -session: Session | None
        -_rewritten_names: dict[str, Path]
        -_must_rewrite: set[str]
        -_writing_pyc: bool
        +find_spec(name, path, target)
        +exec_module(module)
        +mark_rewrite(*names)
    }

    class AssertionRewriter {
        -module_path: str | None
        -config: Config | None
        -source: bytes
        -scope: tuple
        -variables_overwrite: defaultdict
        +run(mod)
        +visit_Assert(assert_)
        +visit_Name(name)
        +visit_BoolOp(boolop)
        +visit_BinOp(binop)
        +visit_Compare(comp)
        +variable()
        +assign(expr)
        +explanation_param(expr)
        +push_format_context()
        +pop_format_context(expl_expr)
    }

    AssertionRewritingHook --> AssertionRewriter : uses
    AssertionRewritingHook ..> Config : depends on
    AssertionRewriter ..> ast.NodeVisitor : inherits

    %% Helper functions are not shown here for brevity

Summary

The [rewrite.py](/projects/286/67427) module is a sophisticated system for enhancing assertion error reporting in pytest by rewriting `assert` statements at import time. It leverages Python's import system hooks and AST transformations to provide detailed introspection into assertion failures, improving developer productivity during test debugging. Through careful integration with pytest's configuration, caching mechanisms, and assertion utilities, it delivers a seamless and performant solution for assertion rewriting.

Appendix: Usage in pytest

When pytest runs tests, it installs the `AssertionRewritingHook` to intercept test module imports. This hook rewrites test modules on the fly so that when an assertion fails, pytest can display comprehensive, user-friendly messages showing the exact values and sub-expressions involved.

Test authors write normal Python `assert` statements as usual; the rewriting and enhanced messages happen transparently.

End of Documentation for rewrite.py