Repository: pytest-dev/pytest
File: src/_pytest/assertion/rewrite.py (Python)

Code before:
            self.enable_assertion_pass_hook = config.getini(
                "enable_assertion_pass_hook"
            )
        else:
            self.enable_assertion_pass_hook = False
        self.source = source
        self.scope: tuple[ast.AST, ...] = ()
        self.variables_overwrite: defaultdict[tuple[ast.AST, ...], dict[str, str]] = (
            defaultdict(dict)
        )

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

        # We'll insert some special imports at the top of the module, but after any
        # docstrings and __future__ imports, so first figure out where that is.
        doc = getattr(mod, "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(doc))) if expect_docstring:
                    if self.is_rewrite_disabled(doc):
                        return
                    expect_docstring = False
                case ast.ImportFrom(level=0, module="__future__"):
                    pass
                case _:
                    break
            pos += 1
        # Special case: for a decorated function, set the lineno to that of the
        # 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
        # Now actually insert the special imports.
        if sys.version_info >= (3, 10):
            aliases = [
                ast.alias("builtins", "@py_builtins", lineno=lineno, col_offset=0),
                ast.alias(
                    "_pytest.assertion.rewrite",
                    "@pytest_ar",
                    lineno=lineno,
                    col_offset=0,
                ),

Review comment (performance):
this looks interesting - whats the performance impact - i would expect match to be faster but that may be wishful thinking

Code after:
            self.enable_assertion_pass_hook = config.getini(
                "enable_assertion_pass_hook"
            )
        else:
            self.enable_assertion_pass_hook = False
        self.source = source
        self.scope: tuple[ast.AST, ...] = ()
        self.variables_overwrite: defaultdict[tuple[ast.AST, ...], dict[str, str]] = (
            defaultdict(dict)
        )

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

        # We'll insert some special imports at the top of the module, but after any
        # docstrings and __future__ imports, so first figure out where that is.
        doc = getattr(mod, "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="__future__"):
                    pass
                case _:
                    break
            pos += 1
        # Special case: for a decorated function, set the lineno to that of the
        # 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
        # Now actually insert the special imports.
        aliases = [
            ast.alias("builtins", "@py_builtins", lineno=lineno, col_offset=0),
            ast.alias(
                "_pytest.assertion.rewrite",
                "@pytest_ar",
                lineno=lineno,
                col_offset=0,