From 9f0d85710681851c04a119f6b0ca742e594a5732 Mon Sep 17 00:00:00 2001 From: Aarav Date: Mon, 19 Oct 2026 04:34:10 +0530 Subject: [PATCH 1/2] Reuse parsed traceback source ASTs within a pytest session Share a bounded source-validated cache between exception reports to avoid repeatedly parsing large modules with many failures. Restore the previous context during config cleanup so nested and failed sessions do not leak their active cache. Co-authored-by: OpenAI Codex --- AUTHORS | 1 + changelog/13360.improvement.rst | 1 + src/_pytest/_code/code.py | 14 +++- src/_pytest/_code/source.py | 39 ++++++++++ src/_pytest/runner.py | 4 + testing/code/test_excinfo.py | 134 ++++++++++++++++++++++++++++++++ testing/code/test_source.py | 33 ++++++++ 7 files changed, 223 insertions(+), 3 deletions(-) create mode 100644 changelog/13360.improvement.rst diff --git a/AUTHORS b/AUTHORS index bd1e27cb37d..c0987a1b882 100644 --- a/AUTHORS +++ b/AUTHORS @@ -3,6 +3,7 @@ merlinux GmbH, Germany, office at merlinux eu Contributors include:: +Aarav Aaron Coleman Abdeali JK Abdelrahman Elbehery diff --git a/changelog/13360.improvement.rst b/changelog/13360.improvement.rst new file mode 100644 index 00000000000..62ff996e37a --- /dev/null +++ b/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. diff --git a/src/_pytest/_code/code.py b/src/_pytest/_code/code.py index f6d2f27672c..57ec010a7b2 100644 --- a/src/_pytest/_code/code.py +++ b/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,17 +290,22 @@ 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 @@ -307,6 +313,8 @@ def getsource( 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] diff --git a/src/_pytest/_code/source.py b/src/_pytest/_code/source.py index 0f1ac00f7b9..ba2c1bedcfb 100644 --- a/src/_pytest/_code/source.py +++ b/src/_pytest/_code/source.py @@ -3,10 +3,13 @@ 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 @@ -14,6 +17,42 @@ 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. diff --git a/src/_pytest/runner.py b/src/_pytest/runner.py index 27c5739845a..611cb2c2d40 100644 --- a/src/_pytest/runner.py +++ b/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: diff --git a/testing/code/test_excinfo.py b/testing/code/test_excinfo.py index 4106a8df671..a91d0e8f849 100644 --- a/testing/code/test_excinfo.py +++ b/testing/code/test_excinfo.py @@ -2296,3 +2296,137 @@ 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='', *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 = "" + 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] + cache: dict[str | Path, ast.AST] = {} + first = entry.getsource(cache) + assert cache + monkeypatch.setattr( + ast, "parse", lambda *args, **kwargs: pytest.fail("reparsed") + ) + assert entry.getsource(cache) == first + finally: + _source_ast_cache.reset(token) diff --git a/testing/code/test_source.py b/testing/code/test_source.py index b2f86024f32..abc458c5adb 100644 --- a/testing/code/test_source.py +++ b/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 From 88586f63b0bc147a12b751c2134907e45e8e11ce Mon Sep 17 00:00:00 2001 From: Aarav Date: Mon, 19 Oct 2026 13:29:54 +0530 Subject: [PATCH 2/2] Cover traceback source lookup without a caller cache --- testing/code/test_excinfo.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/testing/code/test_excinfo.py b/testing/code/test_excinfo.py index a91d0e8f849..30b74e2d6cb 100644 --- a/testing/code/test_excinfo.py +++ b/testing/code/test_excinfo.py @@ -2421,8 +2421,11 @@ def test_traceback_ast_caller_cache_outside_session(monkeypatch: MonkeyPatch) -> 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")