pytest discovery: cache class line numbers to avoid redundant AST parsing

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.<locals>.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"

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("<locals>")
+        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