Method run

src/_pytest/assertion/rewrite.py:683–760

Find all assert statements in *mod* and rewrite them.

(self, mod: ast.Module)

        )

    def run(self, mod: ast.Module) -> None:
        class="st">""class="st">"Find all assert statements in *mod* and rewrite them."class="st">""
        if not mod.body:
            class="cm"># Nothing to do.
            return

        class="cm"># We'll insert some special imports at the top of the module, but after any
        class="cm"># docstrings and __future__ imports, so first figure out where that is.
        doc = getattr(mod, class="st">"docstring", None)
        expect_docstring = doc is None
        if doc is not None and self.is_rewrite_disabled(doc):
            return
        pos = 0
        for item in mod.body:
            match item:
                case ast.Expr(value=ast.Constant(value=str() as doc)) if (
                    expect_docstring
                ):
                    if self.is_rewrite_disabled(doc):
                        return
                    expect_docstring = False
                case ast.ImportFrom(level=0, module=class="st">"__future__"):
                    pass
                case _:
                    break
            pos += 1
        class="cm"># Special case: for a decorated function, set the lineno to that of the
        class="cm"># first decorator, not the `def`. Issue #4984.
        if isinstance(item, ast.FunctionDef) and item.decorator_list:
            lineno = item.decorator_list[0].lineno
        else:
            lineno = item.lineno
        class="cm"># Now actually insert the special imports.
        aliases = [
            ast.alias(class="st">"builtins", class="st">"@py_builtins", lineno=lineno, col_offset=0),
            ast.alias(
                class="st">"_pytest.assertion.rewrite",
                class="st">"@pytest_ar",
                lineno=lineno,
                col_offset=0,
            ),
        ]
        imports = [
            ast.Import([alias], lineno=lineno, col_offset=0) for alias in aliases
        ]
        mod.body[pos:pos] = imports

        class="cm"># Collect asserts.
        self.scope = (mod,)
        nodes: list[ast.AST | Sentinel] = [mod]
        while nodes:
            node = nodes.pop()
            if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef):
                self.scope = tuple((*self.scope, node))
                nodes.append(_SCOPE_END_MARKER)
            if node == _SCOPE_END_MARKER:
                self.scope = self.scope[:-1]
                continue