From 077d5b5712257ea59dc347e0254cd94a4e25bf56 Mon Sep 17 00:00:00 2001 From: Fire162 Date: Fri, 23 Oct 2026 02:48:41 +0530 Subject: [PATCH] pytest discovery: cache class line numbers to avoid redundant AST parsing --- .../tests/pytestadapter/test_utils.py | 90 ++++++++++++++++++- python_files/vscode_pytest/__init__.py | 62 ++++++++++++- 2 files changed, 147 insertions(+), 5 deletions(-) diff --git a/python_files/tests/pytestadapter/test_utils.py b/python_files/tests/pytestadapter/test_utils.py index 70201db7d097..f99cc3e73abd 100644 --- a/python_files/tests/pytestadapter/test_utils.py +++ b/python_files/tests/pytestadapter/test_utils.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. +import ast import os import pathlib import sys @@ -12,7 +13,13 @@ script_dir = pathlib.Path(__file__).parent.parent.parent sys.path.append(os.fspath(script_dir)) -from vscode_pytest import cached_fsdecode, has_symlink_parent # noqa: E402 +from vscode_pytest import ( # noqa: E402 + _CLASS_LINENO_CACHE, + _ClassFinder, + cached_fsdecode, + create_class_node, + has_symlink_parent, +) def test_has_symlink_parent_with_symlink(): @@ -55,3 +62,84 @@ def test_cached_fsdecode(): result3 = cached_fsdecode(test_path2) assert result3 == os.fspath(test_path2) assert result3 != result1 + + +def test_class_finder_extracts_class_lines(): + code = """ +class SimpleClass: + pass + +class OuterClass: + class InnerClass: + pass + +@some_decorator +class DecoratedClass: + pass + +def describe_block(): + def describe_nested(): + pass +""" + tree = ast.parse(code) + finder = _ClassFinder() + finder.visit(tree) + assert finder.classes["SimpleClass"] == 2 + assert finder.classes["OuterClass"] == 5 + assert finder.classes["OuterClass.InnerClass"] == 6 + assert finder.classes["DecoratedClass"] == 9 + assert finder.classes["describe_block"] == 13 + assert finder.classes["describe_block..describe_nested"] == 14 + + +def test_create_class_node_uses_firstlineno(): + class DummyObj: + __firstlineno__ = 42 + + class DummyClassModule: + name = "DummyClass" + nodeid = "test_file.py::DummyClass" + obj = DummyObj() + path = pathlib.Path("test_file.py") + + node = create_class_node(DummyClassModule()) + assert node["lineno"] == "42" + + +def test_create_class_node_caches_ast_by_file(tmp_path): + test_file = tmp_path / "test_sample.py" + test_file.write_text("class TestOne:\n pass\n\nclass TestTwo:\n pass\n") + + class TestOne: + pass + + class TestTwo: + pass + + TestOne.__qualname__ = "TestOne" + TestTwo.__qualname__ = "TestTwo" + + class DummyClass1: + name = "TestOne" + nodeid = f"{test_file}::TestOne" + obj = TestOne + path = test_file + + class DummyClass2: + name = "TestTwo" + nodeid = f"{test_file}::TestTwo" + obj = TestTwo + path = test_file + + # Remove __firstlineno__ if present to exercise AST caching path + if hasattr(TestOne, "__firstlineno__"): + delattr(TestOne, "__firstlineno__") + if hasattr(TestTwo, "__firstlineno__"): + delattr(TestTwo, "__firstlineno__") + + node1 = create_class_node(DummyClass1()) + assert node1["lineno"] == "1" + assert test_file in _CLASS_LINENO_CACHE + + node2 = create_class_node(DummyClass2()) + assert node2["lineno"] == "4" diff --git a/python_files/vscode_pytest/__init__.py b/python_files/vscode_pytest/__init__.py index 88f6327f6e8b..61086265c353 100644 --- a/python_files/vscode_pytest/__init__.py +++ b/python_files/vscode_pytest/__init__.py @@ -3,6 +3,7 @@ from __future__ import annotations +import ast import atexit import contextlib import json @@ -874,6 +875,38 @@ def create_session_node(session: pytest.Session) -> TestNode: } +class _ClassFinder(ast.NodeVisitor): + def __init__(self) -> None: + self.stack: list[str] = [] + self.classes: dict[str, int] = {} + + def visit_FunctionDef(self, node: ast.FunctionDef) -> None: + self.stack.append(node.name) + if node.name.startswith("describe_"): + qualname = ".".join(self.stack) + self.classes[qualname] = ( + node.decorator_list[0].lineno if node.decorator_list else node.lineno + ) + self.stack.append("") + self.generic_visit(node) + self.stack.pop() + self.stack.pop() + + visit_AsyncFunctionDef = visit_FunctionDef # noqa: N815 + + def visit_ClassDef(self, node: ast.ClassDef) -> None: + self.stack.append(node.name) + qualname = ".".join(self.stack) + self.classes[qualname] = ( + node.decorator_list[0].lineno if node.decorator_list else node.lineno + ) + self.generic_visit(node) + self.stack.pop() + + +_CLASS_LINENO_CACHE: dict[pathlib.Path, dict[str, int]] = {} + + def create_class_node(class_module: pytest.Class | DescribeBlockType) -> TestNode: """Creates a class node from a pytest class object. @@ -884,10 +917,31 @@ def create_class_node(class_module: pytest.Class | DescribeBlockType) -> TestNod class_line = "" try: if hasattr(class_module, "obj"): - import inspect - - _, lineno = inspect.getsourcelines(class_module.obj) - class_line = str(lineno) + firstlineno = getattr(class_module.obj, "__firstlineno__", None) + if firstlineno is not None: + class_line = str(firstlineno) + else: + node_path = get_node_path(class_module) + if node_path not in _CLASS_LINENO_CACHE: + try: + tree = ast.parse(node_path.read_text(encoding="utf-8")) + finder = _ClassFinder() + finder.visit(tree) + _CLASS_LINENO_CACHE[node_path] = finder.classes + except Exception: + _CLASS_LINENO_CACHE[node_path] = {} + + qualname = getattr(class_module.obj, "__qualname__", "") + cached_line = _CLASS_LINENO_CACHE[node_path].get(qualname) + if cached_line is None and hasattr(class_module, "name"): + cached_line = _CLASS_LINENO_CACHE[node_path].get(class_module.name) + if cached_line is not None: + class_line = str(cached_line) + else: + import inspect + + _, lineno = inspect.getsourcelines(class_module.obj) + class_line = str(lineno) except (OSError, TypeError): # If we can't get the source lines, leave lineno empty pass