Reuse parsed traceback source ASTs within a session

AUTHORS

@@ -3,6 +3,7 @@ merlinux GmbH, Germany, office at merlinux eu
 
 Contributors include::
 
+Aarav
 Aaron Coleman
 Abdeali JK
 Abdelrahman Elbehery

changelog/13360.improvement.rst

@@ -0,0 +1 @@
+Parsed source ASTs and their statement indexes are reused across traceback reports in the same session, reducing repeated parsing for large test modules with many failures. The cache is bounded and invalidates entries when source changes.

src/_pytest/_code/code.py

@@ -36,6 +36,7 @@
 import pluggy
 
 import _pytest
+from _pytest._code.source import _source_ast_cache
 from _pytest._code.source import findsource
 from _pytest._code.source import getrawcode
 from _pytest._code.source import getstatementrange_ast
@@ -289,24 +290,31 @@ def getsource(
         self, astcache: dict[str | Path, ast.AST] | None = None
     ) -> Source | None:
         """Return failing source code."""
-        # we use the passed in astcache to not reparse asttrees
-        # within exception info printing
+        # Outside a pytest session the caller's cache still shares ASTs within
+        # one exception. A running session also shares them between reports.
         source = self.frame.code.fullsource
         if source is None:
             return None
         key = astnode = None
-        if astcache is not None:
+        session_cache = _source_ast_cache.get()
+        if session_cache is not None:
+            key = self.frame.code.path
+            astnode = session_cache.get(key, source)
+        elif astcache is not None:
             key = self.frame.code.path
             if key is not None:
                 astnode = astcache.get(key, None)
         start = self.getfirstlinesource()
+        cache_miss = astnode is None
         try:
             astnode, _, end = getstatementrange_ast(
                 self.lineno, source, astnode=astnode
             )
         except SyntaxError:
             end = self.lineno + 1
         else:
+            if key is not None and session_cache is not None and cache_miss:
+                session_cache.put(key, source, astnode)
             if key is not None and astcache is not None:
                 astcache[key] = astnode
         return source[start:end]

src/_pytest/_code/source.py

@@ -3,17 +3,56 @@
 
 import ast
 from bisect import bisect_right
+from collections import OrderedDict
 from collections.abc import Iterable
 from collections.abc import Iterator
+from contextvars import ContextVar
 import inspect
 import linecache
+from pathlib import Path
 import textwrap
 import tokenize
 import types
 from typing import overload
 import warnings
 
 
+class _SourceASTCache:
+    """A bounded, source-aware cache shared by a session's traceback reports."""
+
+    def __init__(self, maxsize: int = 100) -> None:
+        self._maxsize = maxsize
+        self._entries: OrderedDict[str | Path, tuple[tuple[str, ...], ast.AST]] = (
+            OrderedDict()
+        )
+
+    def get(self, path: str | Path, source: Source) -> ast.AST | None:
+        cached = self._entries.get(path)
+        if cached is None:
+            return None
+        lines, node = cached
+        # Inspect/linecache may refresh a file or a synthetic source between
+        # failures. Never apply statement locations from a different source.
+        if lines != tuple(source.lines):
+            del self._entries[path]
+            return None
+        self._entries.move_to_end(path)
+        return node
+
+    def put(self, path: str | Path, source: Source, node: ast.AST) -> None:
+        self._entries[path] = (tuple(source.lines), node)
+        self._entries.move_to_end(path)
+        if len(self._entries) > self._maxsize:
+            self._entries.popitem(last=False)
+
+
+# No process-global cache: each session installs its own and restores the
+# previous value on cleanup, including nested pytest.main() invocations.
+_source_ast_cache: ContextVar[_SourceASTCache | None] = ContextVar(
+    "pytest_source_ast_cache", default=None
+)
+
+
 class Source:
     """An immutable object holding a source code fragment.
 

src/_pytest/runner.py

@@ -25,6 +25,8 @@
 from _pytest._code.code import ExceptionChainRepr
 from _pytest._code.code import ExceptionInfo
 from _pytest._code.code import TerminalRepr
+from _pytest._code.source import _source_ast_cache
+from _pytest._code.source import _SourceASTCache
 from _pytest.config.argparsing import Parser
 from _pytest.deprecated import check_ispytest
 from _pytest.nodes import Collector
@@ -106,6 +108,8 @@ def pytest_terminal_summary(terminalreporter: TerminalReporter) -> None:
 
 def pytest_sessionstart(session: Session) -> None:
     session._setupstate = SetupState()
+    token = _source_ast_cache.set(_SourceASTCache())
+    session.config.add_cleanup(lambda: _source_ast_cache.reset(token))
 
 
 def pytest_sessionfinish(session: Session) -> None:

testing/code/test_excinfo.py

@@ -2296,3 +2296,140 @@ def test_check_error_notes_failure(
     with pytest.raises(AssertionError):
         with pytest.raises(type(error), match=match):
             raise error
+
+
+def test_traceback_ast_reused_between_reports(pytester: Pytester) -> None:
+    """Parsing the whole module once per failure is quadratic (#13360)."""
+    pytester.makeconftest(
+        """
+        import ast
+        from unittest.mock import patch
+
+        parses = 0
+        original_parse = ast.parse
+
+        def parse(source, filename='<unknown>', *args, **kwargs):
+            global parses
+            if filename == 'source' and 'TRACEBACK_AST_CACHE_SENTINEL' in source:
+                parses += 1
+            return original_parse(source, filename, *args, **kwargs)
+
+        def pytest_sessionstart(session):
+            patcher = patch.object(ast, 'parse', parse)
+            patcher.start()
+            session.config.add_cleanup(patcher.stop)
+
+        def pytest_sessionfinish():
+            print(f'TRACEBACK_AST_PARSES={parses}')
+        """
+    )
+    pytester.makepyfile(
+        "# TRACEBACK_AST_CACHE_SENTINEL\n"
+        + "\n".join(
+            f'def test_{i}():\n    assert False, "failure {i}"' for i in range(10)
+        )
+    )
+    result = pytester.runpytest("--tb=short")
+    result.assert_outcomes(failed=10)
+    result.stdout.fnmatch_lines(["*TRACEBACK_AST_PARSES=1"])
+    for i in range(10):
+        result.stdout.fnmatch_lines([f"*AssertionError: failure {i}*"])
+
+
+def test_traceback_ast_cache_refreshes_synthetic_source(
+    monkeypatch: MonkeyPatch,
+) -> None:
+    import linecache
+
+    from _pytest._code.source import _source_ast_cache
+    from _pytest._code.source import _SourceASTCache
+
+    filename = "<changing traceback source>"
+    cache = _SourceASTCache()
+    token = _source_ast_cache.set(cache)
+    try:
+        # The second statement has a different end position. Reusing the first
+        # AST would omit the closing parenthesis from the second traceback.
+        for source in (
+            "def fail():\n    raise ValueError('first')\n",
+            "def fail():\n    raise ValueError(\n        'second'\n    )\n",
+        ):
+            lines = source.splitlines(keepends=True)
+            monkeypatch.setitem(
+                linecache.cache, filename, (len(source), None, lines, filename)
+            )
+            namespace: dict[str, Any] = {}
+            exec(compile(source, filename, "exec"), namespace)
+            with pytest.raises(ValueError) as excinfo:
+                namespace["fail"]()
+            entry = excinfo.traceback[-1]
+            assert str(entry.getsource()) == source.rstrip("\n")
+            # Mutating a returned Source must not corrupt the next lookup.
+            extracted = entry.getsource()
+            assert extracted is not None
+            extracted.lines.clear()
+            assert str(entry.getsource()) == source.rstrip("\n")
+    finally:
+        _source_ast_cache.reset(token)
+
+
+def test_traceback_ast_cache_session_lifetime(pytester: Pytester) -> None:
+    from _pytest._code.source import _source_ast_cache
+    from _pytest._code.source import _SourceASTCache
+
+    previous = _source_ast_cache.get()
+    seen: list[_SourceASTCache | None] = []
+
+    class Plugin:
+        @pytest.hookimpl(trylast=True)
+        def pytest_sessionstart(self) -> None:
+            seen.append(_source_ast_cache.get())
+
+    pytester.makepyfile("def test_pass(): pass")
+    for _ in range(2):
+        pytester.inline_run(plugins=[Plugin()]).assertoutcome(passed=1)
+        assert _source_ast_cache.get() is previous
+    assert all(cache is not None and cache is not previous for cache in seen)
+    assert seen[0] is not seen[1]
+
+
+def test_traceback_ast_cache_failed_session_cleanup(pytester: Pytester) -> None:
+    from _pytest._code.source import _source_ast_cache
+
+    previous = _source_ast_cache.get()
+    pytester.makeconftest(
+        """
+        import pytest
+
+        @pytest.hookimpl(trylast=True)
+        def pytest_sessionstart():
+            raise RuntimeError('session start failed')
+        """
+    )
+    result = pytester.runpytest()
+    assert result.ret == pytest.ExitCode.INTERNAL_ERROR
+    assert _source_ast_cache.get() is previous
+
+
+def test_traceback_ast_caller_cache_outside_session(monkeypatch: MonkeyPatch) -> None:
+    import ast
+
+    from _pytest._code.source import _source_ast_cache
+
+    token = _source_ast_cache.set(None)
+    try:
+        with pytest.raises(ValueError) as excinfo:
+            raise ValueError("caller owns the cache")
+        entry = excinfo.traceback[-1]
+        uncached = entry.getsource()
+        assert uncached is not None
+        cache: dict[str | Path, ast.AST] = {}
+        first = entry.getsource(cache)
+        assert first == uncached
+        assert cache
+        monkeypatch.setattr(
+            ast, "parse", lambda *args, **kwargs: pytest.fail("reparsed")
+        )
+        assert entry.getsource(cache) == first
+    finally:
+        _source_ast_cache.reset(token)

testing/code/test_source.py

@@ -724,3 +724,36 @@ def test_statement_linenos_are_memoized_per_node() -> None:
     other = ast.parse("z = 3\n")
     assert _statement_linenos(other) == [0]
     assert _statement_linenos(node) is values
+
+
+def test_session_source_ast_cache_source_changes() -> None:
+    from _pytest._code.source import _SourceASTCache
+
+    cache = _SourceASTCache()
+    source = Source("first = 1")
+    first = ast.parse(str(source))
+    cache.put("test.py", source, first)
+    assert cache.get("test.py", Source("first = 1")) is first
+    assert cache.get("other.py", source) is None
+    # The snapshot must not change when a caller mutates the original Source.
+    source.lines[:] = ["second = 2", "third = 3"]
+    assert cache.get("test.py", source) is None
+    second = ast.parse(str(source))
+    cache.put("test.py", source, second)
+    assert cache.get("test.py", source) is second
+    assert cache.get("test.py", Source("first = 1")) is None
+
+
+def test_session_source_ast_cache_eviction() -> None:
+    from _pytest._code.source import _SourceASTCache
+
+    cache = _SourceASTCache(maxsize=2)
+    source = Source("value = 1")
+    node = ast.parse(str(source))
+    cache.put("first.py", source, node)
+    cache.put("second.py", source, node)
+    assert cache.get("first.py", source) is node
+    cache.put("third.py", source, node)
+    assert cache.get("second.py", source) is None
+    assert cache.get("first.py", source) is node
+    assert cache.get("third.py", source) is node