From 64d7852a290231a61f111c3021052076f3fdd73f Mon Sep 17 00:00:00 2001 From: ProtocolWarden <32967198+ProtocolWarden@users.noreply.github.com> Date: Tue, 14 Jul 2026 14:15:59 -0400 Subject: [PATCH] fix(Custodian): reuse AST caches for watchdog audit path --- src/custodian/audit_kit/detectors/docs.py | 140 +++++++++++++----- src/custodian/audit_kit/detectors/stubs.py | 101 +++++++------ .../audit_kit/detectors/test_shape.py | 44 +++--- tests/test_test_shape_detectors.py | 6 +- 4 files changed, 191 insertions(+), 100 deletions(-) diff --git a/src/custodian/audit_kit/detectors/docs.py b/src/custodian/audit_kit/detectors/docs.py index db5575a..fcdab03 100644 --- a/src/custodian/audit_kit/detectors/docs.py +++ b/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): diff --git a/src/custodian/audit_kit/detectors/stubs.py b/src/custodian/audit_kit/detectors/stubs.py index ebe52fd..f6f3d05 100644 --- a/src/custodian/audit_kit/detectors/stubs.py +++ b/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 diff --git a/src/custodian/audit_kit/detectors/test_shape.py b/src/custodian/audit_kit/detectors/test_shape.py index 96ecbf5..5a23e46 100644 --- a/src/custodian/audit_kit/detectors/test_shape.py +++ b/src/custodian/audit_kit/detectors/test_shape.py @@ -71,6 +71,7 @@ _MAX_SAMPLES = 8 _NEEDS_TF = frozenset({"ast_forest", "tests_forest"}) +_NEEDS_TESTS_ONLY = frozenset({"tests_forest"}) def build_test_shape_detectors() -> list[Detector]: @@ -78,19 +79,19 @@ def build_test_shape_detectors() -> list[Detector]: Detector("T1", "public src symbol with no reference in tests", "open", detect_t1, LOW, _NEEDS_TF), Detector("T2", "test function with no assert statement", "open", - detect_t2, LOW), + detect_t2, LOW, _NEEDS_TESTS_ONLY), Detector("T3", "unconditional pytest.skip without environment gate", "open", - detect_t3, LOW), + detect_t3, LOW, _NEEDS_TESTS_ONLY), Detector("T4", "pytest fixture defined but never requested by any test or fixture", "open", - detect_t4, LOW), + detect_t4, LOW, _NEEDS_TESTS_ONLY), Detector("T5", "pytest.mark.parametrize with a single test case — should be a plain test", "open", - detect_t5, LOW), + detect_t5, LOW, _NEEDS_TESTS_ONLY), Detector("T6", "src module is never imported by any test file", "open", detect_t6, LOW, _NEEDS_TF), Detector("T7", "src module has no parallel test file under tests/", "open", detect_t7, LOW), Detector("T8", "test file imports nothing from any src package", "open", - detect_t8, LOW), + detect_t8, LOW, _NEEDS_TESTS_ONLY), ] @@ -167,6 +168,12 @@ def _parse_test_files(tests_root: Path) -> list[tuple[Path, ast.Module]]: return results +def _iter_test_files(context: AuditContext) -> list[tuple[Path, ast.Module, str]]: + if context.graph is not None and context.graph.tests_forest is not None: + return list(context.graph.tests_forest.items()) + return [(path, tree, "") for path, tree in _parse_test_files(context.tests_root)] + + # ── T1 ──────────────────────────────────────────────────────────────────────── def _t1_excluded_paths(context: AuditContext) -> set[str]: @@ -237,7 +244,7 @@ def detect_t2(context: AuditContext) -> DetectorResult: samples: list[str] = [] count = 0 - for path, tree in _parse_test_files(context.tests_root): + for path, tree, _src in _iter_test_files(context): rel = path.relative_to(context.repo_root) rel_posix = rel.as_posix() if t2_excludes and any(glob_match(rel_posix, excl) for excl in t2_excludes): @@ -278,12 +285,9 @@ def detect_t3(context: AuditContext) -> DetectorResult: samples: list[str] = [] count = 0 - for path, _tree in _parse_test_files(context.tests_root): + for path, _tree, src in _iter_test_files(context): rel = path.relative_to(context.repo_root) - try: - lines = path.read_text(encoding="utf-8").splitlines() - except OSError: - continue + lines = src.splitlines() if src else path.read_text(encoding="utf-8").splitlines() for i, line in enumerate(lines, 1): stripped = line.lstrip() is_call = "pytest.skip(" in line @@ -346,9 +350,9 @@ def detect_t4(context: AuditContext) -> DetectorResult: # Pass 2: collect all parameter names across test functions and fixtures requested_names: set[str] = set() - all_files: list[tuple[Path, ast.Module]] = _parse_test_files(context.tests_root) + all_files = _iter_test_files(context) - for path, tree in all_files: + for path, tree, _src in all_files: rel_str = str(path.relative_to(context.repo_root)) if globs and any(PurePosixPath(rel_str).match(g) for g in globs): continue @@ -427,7 +431,7 @@ def detect_t5(context: AuditContext) -> DetectorResult: count = 0 samples: list[str] = [] - for path, tree in _parse_test_files(context.tests_root): + for path, tree, _src in _iter_test_files(context): rel = str(path.relative_to(context.repo_root)) for node in ast.walk(tree): if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): @@ -672,7 +676,7 @@ def _file_touches_src(tree: ast.AST, src_packages: set[str]) -> bool: def _conftest_dirs_touching_src( - tests_root: Path, src_packages: set[str], + context: AuditContext, src_packages: set[str], ) -> set[Path]: """Return dir paths whose conftest.py (or any ancestor's conftest.py) imports a src package. @@ -682,10 +686,8 @@ def _conftest_dirs_touching_src( tests under it implicitly exercise src — they're not dangling. """ touching: set[Path] = set() - for conftest in tests_root.rglob("conftest.py"): - try: - tree = ast.parse(conftest.read_text(encoding="utf-8", errors="replace")) - except (OSError, SyntaxError): + for conftest, tree, _src in _iter_test_files(context): + if conftest.name != "conftest.py": continue if _file_touches_src(tree, src_packages): touching.add(conftest.parent.resolve()) @@ -731,11 +733,11 @@ def detect_t8(context: AuditContext) -> DetectorResult: if not src_packages: return DetectorResult(count=0, samples=[]) - conftest_dirs = _conftest_dirs_touching_src(context.tests_root, src_packages) + conftest_dirs = _conftest_dirs_touching_src(context, src_packages) samples: list[str] = [] count = 0 - for path, tree in _parse_test_files(context.tests_root): + for path, tree, _src in _iter_test_files(context): if path.name == "conftest.py" or path.name == "__init__.py": continue rel = path.relative_to(context.repo_root) diff --git a/tests/test_test_shape_detectors.py b/tests/test_test_shape_detectors.py index fca376b..4eecd13 100644 --- a/tests/test_test_shape_detectors.py +++ b/tests/test_test_shape_detectors.py @@ -9,6 +9,7 @@ from custodian.audit_kit.detector import AnalysisGraph, AuditContext from custodian.audit_kit.detectors.test_shape import detect_t2 from custodian.audit_kit.passes.ast_forest import AstForest +from custodian.audit_kit.passes.tests_forest import build_tests_forest def _write_test_file(src: str, tmp_path: Path, name: str = "test_example.py") -> None: @@ -19,13 +20,14 @@ def _write_test_file(src: str, tmp_path: Path, name: str = "test_example.py") -> def _ctx(tmp_path: Path) -> AuditContext: (tmp_path / "src").mkdir(parents=True, exist_ok=True) + tests_root = tmp_path / "tests" return AuditContext( repo_root=tmp_path, src_root=tmp_path / "src", - tests_root=tmp_path / "tests", + tests_root=tests_root, config={}, plugin_modules=[], - graph=AnalysisGraph(ast_forest=AstForest()), + graph=AnalysisGraph(ast_forest=AstForest(), tests_forest=build_tests_forest(tests_root)), )