fix(Custodian): reuse AST caches for watchdog audit path

src/custodian/audit_kit/detectors/docs.py

@@ -39,6 +39,7 @@
 from __future__ import annotations
 
 import ast as _ast
+from dataclasses import dataclass
 import re
 from pathlib import Path
 
@@ -48,18 +49,30 @@
 from custodian.audit_kit.code_health import _py_files
 
 _MAX_SAMPLES = 8
+_NEEDS_AST = frozenset({"ast_forest"})
+_NEEDS_BOTH = frozenset({"ast_forest", "tests_forest"})
+_DOC_INDEX_CACHE: dict[tuple[int, int], "_DocIndex"] = {}
+
+
+@dataclass(frozen=True)
+class _DocIndex:
+    defined_symbols: set[str]
+    field_symbols: set[str]
+    event_symbols: set[str]
+    string_literals: set[str]
+    tests_defined_symbols: set[str]
 
 
 def build_docs_detectors() -> list[Detector]:
     return [
         Detector("K1", "doc references a symbol not found in src (phantom symbol)", "open",
-                 detect_k1, LOW),
+                 detect_k1, LOW, _NEEDS_BOTH),
         Detector("K2", "doc cites a value not present as string literal in src (value drift)", "open",
-                 detect_k2, LOW),
+                 detect_k2, LOW, _NEEDS_AST),
         Detector("K3", "docstring Args section names parameter not in function signature (param drift)", "open",
-                 detect_k3, LOW),
+                 detect_k3, LOW, _NEEDS_AST),
         Detector("K4", "docstring Args type does not match signature annotation (type drift)", "open",
-                 detect_k4, LOW),
+                 detect_k4, LOW, _NEEDS_AST),
     ]
 
 
@@ -117,6 +130,57 @@ def _build_src_text(context: AuditContext) -> tuple[str, str]:
     return src_text, tests_text
 
 
+def _doc_index(context: AuditContext) -> _DocIndex:
+    graph = context.graph
+    ast_forest = None if graph is None else graph.ast_forest
+    tests_forest = None if graph is None else graph.tests_forest
+    cache_key = (id(ast_forest), id(tests_forest))
+    cached = _DOC_INDEX_CACHE.get(cache_key)
+    if cached is not None:
+        return cached
+
+    defined_symbols: set[str] = set()
+    field_symbols: set[str] = set()
+    event_symbols: set[str] = set()
+    string_literals: set[str] = set()
+    tests_defined_symbols: set[str] = set()
+
+    if ast_forest is not None:
+        for _path, tree, _src in ast_forest.items():
+            for node in _ast.walk(tree):
+                if isinstance(node, (_ast.FunctionDef, _ast.AsyncFunctionDef, _ast.ClassDef)):
+                    defined_symbols.add(node.name)
+                elif isinstance(node, _ast.AnnAssign) and isinstance(node.target, _ast.Name):
+                    field_symbols.add(node.target.id)
+                elif isinstance(node, _ast.Constant) and isinstance(node.value, str):
+                    string_literals.add(node.value)
+                elif isinstance(node, _ast.Dict):
+                    for key, value in zip(node.keys, node.values, strict=False):
+                        if (
+                            isinstance(key, _ast.Constant)
+                            and key.value == "event"
+                            and isinstance(value, _ast.Constant)
+                            and isinstance(value.value, str)
+                        ):
+                            event_symbols.add(value.value)
+
+    if tests_forest is not None:
+        for _path, tree, _src in tests_forest.items():
+            for node in _ast.walk(tree):
+                if isinstance(node, (_ast.FunctionDef, _ast.AsyncFunctionDef, _ast.ClassDef)):
+                    tests_defined_symbols.add(node.name)
+
+    index = _DocIndex(
+        defined_symbols=defined_symbols,
+        field_symbols=field_symbols,
+        event_symbols=event_symbols,
+        string_literals=string_literals,
+        tests_defined_symbols=tests_defined_symbols,
+    )
+    _DOC_INDEX_CACHE[cache_key] = index
+    return index
+
+
 _DEFERRED_WORDS = ("deferred", "out of scope", "not yet implemented", "future:", "deprecated")
 _IMPL_MARKER_RE = re.compile(
     r"\*\*Files:\*\*|\bImplementation:|see\s+`|defined in `|"
@@ -139,21 +203,17 @@ def detect_k1(context: AuditContext) -> DetectorResult:
     common_words: set[str] = set(audit_cfg.get("common_words") or [])
     stale_handlers: set[str] = set(audit_cfg.get("stale_handlers") or [])
 
-    src_text, tests_text = _build_src_text(context)
+    index = _doc_index(context)
 
     def _exists(name: str) -> bool:
         if name in common_words or name in stale_handlers:
             return True
-        if re.search(rf"\b(def|class)\s+{re.escape(name)}\b", src_text):
-            return True
-        if re.search(rf"^\s+{re.escape(name)}\s*:\s*[A-Za-z]", src_text, re.MULTILINE):
-            return True
-        if re.search(rf"\b(def|class)\s+{re.escape(name)}\b", tests_text):
-            return True
-        # Exists as a quoted string literal in src (e.g. dict key, enum value, config key)
-        if re.search(rf"""['"]{re.escape(name)}['"]""", src_text):
-            return True
-        return False
+        return (
+            name in index.defined_symbols
+            or name in index.field_symbols
+            or name in index.tests_defined_symbols
+            or name in index.string_literals
+        )
 
     seen: dict[str, tuple[Path, int]] = {}
     for f in _doc_files(context.repo_root, audit_cfg):
@@ -212,7 +272,7 @@ def detect_k2(context: AuditContext) -> DetectorResult:
     extra_known: set[str] = {v.lower() for v in (audit_cfg.get("known_values") or [])}
     known_values = _DEFAULT_KNOWN_VALUES | extra_known
 
-    src_text, _ = _build_src_text(context)
+    index = _doc_index(context)
 
     seen: dict[str, tuple[Path, int]] = {}
     for f in _doc_files(context.repo_root, audit_cfg):
@@ -227,11 +287,11 @@ def detect_k2(context: AuditContext) -> DetectorResult:
                 name = m.group(1)
                 if name in seen or name.lower() in known_values:
                     continue
-                if re.search(rf"""['"]{re.escape(name)}['"]""", src_text):
-                    continue
-                if re.search(rf"^\s+{re.escape(name)}\s*:\s*[A-Za-z]", src_text, re.MULTILINE):
-                    continue
-                if re.search(rf"\b(def|class)\s+{re.escape(name)}\b", src_text):
+                if (
+                    name in index.string_literals
+                    or name in index.field_symbols
+                    or name in index.defined_symbols
+                ):
                     continue
                 seen[name] = (f, i)
 
@@ -346,17 +406,24 @@ def detect_k3(context: AuditContext) -> DetectorResult:
     samples: list[str] = []
     count = 0
 
-    for path in _py_files(context, "K3"):
+    if context.graph is None or context.graph.ast_forest is None:
+        file_iter = []
+        for path in _py_files(context, "K3"):
+            try:
+                raw = path.read_text(encoding="utf-8")
+                tree = _ast.parse(raw)
+            except (OSError, UnicodeDecodeError, SyntaxError):
+                continue
+            file_iter.append((path, tree))
+    else:
+        file_iter = [(path, tree) for path, tree, _src in context.graph.ast_forest.items()]
+
+    for path, tree in file_iter:
         if globs:
             from custodian.audit_kit.code_health import _glob_to_regex
             rel_str = str(path.relative_to(context.repo_root))
             if any(_glob_to_regex(g).match(rel_str) for g in globs):
                 continue
-        try:
-            raw = path.read_text(encoding="utf-8")
-            tree = _ast.parse(raw)
-        except (OSError, UnicodeDecodeError, SyntaxError):
-            continue
         rel = path.relative_to(context.repo_root)
 
         for node in _ast.walk(tree):
@@ -524,17 +591,24 @@ def detect_k4(context: AuditContext) -> DetectorResult:
     samples: list[str] = []
     count = 0
 
-    for path in _py_files(context, "K4"):
+    if context.graph is None or context.graph.ast_forest is None:
+        file_iter = []
+        for path in _py_files(context, "K4"):
+            try:
+                raw = path.read_text(encoding="utf-8")
+                tree = _ast.parse(raw)
+            except (OSError, UnicodeDecodeError, SyntaxError):
+                continue
+            file_iter.append((path, tree))
+    else:
+        file_iter = [(path, tree) for path, tree, _src in context.graph.ast_forest.items()]
+
+    for path, tree in file_iter:
         if globs:
             from custodian.audit_kit.code_health import _glob_to_regex
             rel_str = str(path.relative_to(context.repo_root))
             if any(_glob_to_regex(g).match(rel_str) for g in globs):
                 continue
-        try:
-            raw = path.read_text(encoding="utf-8")
-            tree = _ast.parse(raw)
-        except (OSError, UnicodeDecodeError, SyntaxError):
-            continue
         rel = path.relative_to(context.repo_root)
 
         for node in _ast.walk(tree):

src/custodian/audit_kit/detectors/stubs.py

@@ -44,6 +44,7 @@ class itself (not inherited from further bases).  Excludes abstract
 from __future__ import annotations
 
 import ast
+from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -56,6 +57,20 @@ class itself (not inherited from further bases).  Excludes abstract
 
 _MAX_SAMPLES = 8
 _NEEDS = frozenset({"ast_forest"})
+_STUB_SCAN_CACHE: dict[int, "_StubFileInfo"] = {}
+
+
+@dataclass(frozen=True)
+class _StubFunctionInfo:
+    node: ast.FunctionDef | ast.AsyncFunctionDef
+    container: ast.ClassDef | None
+
+
+@dataclass(frozen=True)
+class _StubFileInfo:
+    protocol_names: set[str]
+    except_fn_ids: set[int]
+    functions: list[_StubFunctionInfo]
 
 
 def build_stub_detectors() -> list[Detector]:
@@ -119,49 +134,49 @@ def _is_ellipsis_only(stmt: ast.stmt) -> bool:
 
 
 def _protocol_classes(tree: ast.Module) -> set[str]:
-    """Return names of Protocol-subclassing classes in this module."""
-    names: set[str] = set()
-    for node in ast.walk(tree):
-        if not isinstance(node, ast.ClassDef):
-            continue
-        for base in node.bases:
-            base_name = None
-            if isinstance(base, ast.Name):
-                base_name = base.id
-            elif isinstance(base, ast.Attribute):
-                base_name = base.attr
-            if base_name == "Protocol":
-                names.add(node.name)
-    return names
-
+    return _stub_file_info(tree).protocol_names
 
-def _except_handler_functions(tree: ast.Module) -> set[int]:
-    """Return ids of FunctionDef nodes that live inside except-handler bodies.
 
-    try/except fallback stubs (e.g. ``except ImportError: class Foo: def add():...``)
-    are intentional no-ops, not unfinished implementations.
-    """
-    ids: set[int] = set()
-    for node in ast.walk(tree):
-        if not isinstance(node, ast.ExceptHandler):
-            continue
-        for child in ast.walk(ast.Module(body=node.body, type_ignores=[])):
-            if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
-                ids.add(id(child))
-    return ids
+def _stub_file_info(tree: ast.Module) -> _StubFileInfo:
+    cached = _STUB_SCAN_CACHE.get(id(tree))
+    if cached is not None:
+        return cached
 
+    protocol_names: set[str] = set()
+    except_fn_ids: set[int] = set()
+    functions: list[_StubFunctionInfo] = []
+    class_by_child_id: dict[int, ast.ClassDef] = {}
 
-def _containing_class(
-    func: ast.FunctionDef | ast.AsyncFunctionDef,
-    tree: ast.Module,
-) -> ast.ClassDef | None:
-    """Return the ClassDef that directly contains func, or None."""
     for node in ast.walk(tree):
         if isinstance(node, ast.ClassDef):
+            for base in node.bases:
+                base_name = None
+                if isinstance(base, ast.Name):
+                    base_name = base.id
+                elif isinstance(base, ast.Attribute):
+                    base_name = base.attr
+                if base_name == "Protocol":
+                    protocol_names.add(node.name)
             for item in node.body:
-                if item is func:
-                    return node
-    return None
+                if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                    class_by_child_id[id(item)] = node
+        elif isinstance(node, ast.ExceptHandler):
+            for child in ast.walk(ast.Module(body=node.body, type_ignores=[])):
+                if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                    except_fn_ids.add(id(child))
+        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            functions.append(_StubFunctionInfo(node=node, container=None))
+
+    info = _StubFileInfo(
+        protocol_names=protocol_names,
+        except_fn_ids=except_fn_ids,
+        functions=[
+            _StubFunctionInfo(node=fn.node, container=class_by_child_id.get(id(fn.node)))
+            for fn in functions
+        ],
+    )
+    _STUB_SCAN_CACHE[id(tree)] = info
+    return info
 
 
 def _sample(
@@ -201,18 +216,16 @@ def _scan_functions(
     for path, tree, _src in context.graph.ast_forest.items():
         if str(path) in excluded_paths:
             continue
-        protocol_names = _protocol_classes(tree)
-        except_fn_ids = _except_handler_functions(tree)
+        file_info = _stub_file_info(tree)
 
-        for node in ast.walk(tree):
-            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
-                continue
+        for function_info in file_info.functions:
+            node = function_info.node
             if _has_decorator(node, "abstractmethod", "overload"):
                 continue
-            if id(node) in except_fn_ids:
+            if id(node) in file_info.except_fn_ids:
                 continue
-            container = _containing_class(node, tree)
-            if container and container.name in protocol_names:
+            container = function_info.container
+            if container and container.name in file_info.protocol_names:
                 continue
             if predicate(node, container):
                 count += 1