scripts.extract_reference_edges

Reference › scripts › extract_reference_edges

Extracts non-import reference edges from Python source into a JSON map: references that an import graph does not show, such as a script naming another file's path or launching another script. affected_tests regenerates this index while selecting tests and reads it through its ReferenceIndex. tests/scripts/test_extract_reference_edges.py and tests/scripts/test_extract_reference_edges_coverage.py import it, and it writes through atomic_json.

Invocation: main requires --source (one or more repo-relative source roots) and --output (the JSON file). It prints the edge count to standard error and returns 0. The output holds a version number and the list of edges.

extract_edges walks the source roots, optionally in parallel, and _extract_file parses one file with a _RefEdgeVisitor. The visitor first collects constants that hold path fragments (collect_path_consts), then records edges from calls, from argument vectors passed to subprocess calls (_argv_paths), from path-building expressions such as a division of paths (_div_target, _path_segments) and from string constants that look like repository paths (_is_repo_path, visit_Constant). A map of basenames under the Claude configuration directory lets bare names resolve (_claude_basename_map). The output is written atomically because a reader may have the file mapped while it is being regenerated.

At a glance

Class diagram

Dependency diagram

API

Classes

_RefEdgeVisitor

Bases: ast.NodeVisitor · Line 156

Harvest reference edges from one module's call nodes.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/scripts/extract_reference_edges.py

#!/usr/bin/env python3
"""Extract NON-IMPORT reference edges from Python source into a JSON map.

The affected-test resolver (``scripts/affected_tests.py``) walks the IMPORT
closure — which is exact for Python imports and blind everywhere else. The
three edge classes that actually burned full-suite runs (2026-08-11: a
2-file deploy/hook change escalated to 32k tests) are statically visible
but are not imports:

- **invoke** — ``subprocess.run(["bash", ".../hook.sh"])`` and friends:
  a script path sitting in an argv literal. This is how hook tests bind to
  ``.claude/hooks/*.sh`` and how CLIs invoke sibling scripts.
- **bdd** — ``scenario("x.feature", ...)`` / ``scenarios(...)``: pytest-bdd
  binds a feature file by string, not by import.
- **consult** — Prolog rule files loaded via ``consult("...pl")``-style
  calls at runtime.

This extractor walks call nodes (stdlib ``ast`` — no import graph needed,
each file parses independently) and harvests repo-relative string literals
in those positions. What it cannot see stays unseen BY DESIGN: a
runtime-computed path (f-string, variable, os.path.join) yields no edge,
and the resolver treats an unreferenced non-documentation file as UNKNOWN
(full-suite escalation) — fail-closed, never a silent skip.

Output (deterministic — sorted, no timestamps):

    {"version": 1,
     "edges": [{"from": "tests/scripts/test_x.py",
                "kind": "invoke"|"bdd"|"consult",
                "target": ".claude/hooks/x.sh"}, ...]}

CLI:
    python scripts/extract_reference_edges.py \
        --source src scripts tests --output plans/reference-edges.json
"""

from __future__ import annotations

import argparse
import ast
import functools
import os
import sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

#: Argv head names that mark a subprocess call's list literal as a shell
#: invocation of a script.
SHELL_NAMES = frozenset({"bash", "sh", "zsh", "dash"})

#: subprocess entry points whose first positional arg is an argv list.
SUBPROCESS_FUNCS = frozenset({"run", "Popen", "check_output", "check_call", "call"})

#: pytest-bdd binding functions whose string args name feature files.
BDD_FUNCS = frozenset({"scenario", "scenarios"})

#: Content-read entry points whose first arg is a file being asserted on.
READ_FUNCS = frozenset({"read_text", "read_bytes", "open"})

#: Suffixes that make an argv-literal string a repo-artifact reference.
PATH_SUFFIXES = (".sh", ".py", ".pl", ".feature", ".json")


def _dotted(func: ast.expr) -> str:
    """Best-effort dotted name of a call target (``subprocess.run`` etc.)."""
    parts: list[str] = []
    node = func
    while isinstance(node, ast.Attribute):
        parts.append(node.attr)
        node = node.value
    if isinstance(node, ast.Name):
        parts.append(node.id)
    return ".".join(reversed(parts))


def _last(name: str) -> str:
    return name.rsplit(".", 1)[-1] if name else ""


def _is_repo_path(s: str) -> bool:
    """A literal that names a repo artifact, not a flag/URL/abs path."""
    if not s.endswith(PATH_SUFFIXES):
        return False
    if (
        s.startswith(("/", "http://", "https://"))
        or "${" in s
        or "/" not in s.lstrip("./")
    ):
        return False
    return True


def _normalize(s: str) -> str:
    return s[2:] if s.startswith("./") else s


#: Variable names accepted as repo-root anchors for module-level path
#: constants. Uppercase-only by convention: fixture parameters like
#: ``repo_root``/``tmp_path`` (lowercase) point at TEMP trees whose
#: ``src/...`` segments would otherwise mint edges to real repo files they
#: merely shadow — a false edge is over-selection (safe) but a false edge
#: that suppresses escalation is not, so the anchor set is deliberately
#: narrow (2026-08-11 design review).
ROOT_NAMES = frozenset(
    {"REPO_ROOT", "ROOT", "PROJECT_ROOT", "PACKAGE_ROOT", "_REPO_ROOT", "_ROOT"}
)


def _path_segments(expr: ast.expr, consts: dict[str, str]) -> list[str] | None:
    """Resolve a ``BASE / "a" / "b.sh"`` or ``BASE.joinpath(...)`` chain to
    its string segments. Returns None for anything not statically rooted in
    a ROOT_NAMES variable (directly, or via another module-level constant).
    """
    if isinstance(expr, ast.BinOp) and isinstance(expr.op, ast.Div):
        left = _path_segments(expr.left, consts)
        if left is None:
            return None
        if isinstance(expr.right, ast.Constant) and isinstance(expr.right.value, str):
            return left + [expr.right.value]
        return None
    if isinstance(expr, ast.Call):
        if isinstance(expr.func, ast.Attribute) and expr.func.attr == "joinpath":
            base = _path_segments(expr.func.value, consts)
            if base is None:
                return None
            segs = [
                a.value
                for a in expr.args
                if isinstance(a, ast.Constant) and isinstance(a.value, str)
            ]
            if len(segs) != len(expr.args):
                return None
            return base + segs
        # str(<path expr>) — unwrap.
        if (
            isinstance(expr.func, ast.Name)
            and expr.func.id == "str"
            and len(expr.args) == 1
        ):
            return _path_segments(expr.args[0], consts)
        return None
    if isinstance(expr, ast.Name):
        if expr.id in ROOT_NAMES:
            return []
        if expr.id in consts:
            resolved = _path_segments_str(consts[expr.id])
            return resolved
        return None
    return None


def _path_segments_str(joined: str) -> list[str]:
    return joined.split("/") if joined else []


class _RefEdgeVisitor(ast.NodeVisitor):
    """Harvest reference edges from one module's call nodes."""

    def __init__(
        self, rel_path: str, *, claude_basenames: dict[str, list[str]] | None = None
    ) -> None:
        self.rel_path = rel_path
        self.edges: list[dict[str, str]] = []
        # Module-level path constants (``HOOK_PATH = REPO_ROOT / ...``),
        # filled by ``collect_path_consts`` before the call walk.
        self.path_consts: dict[str, str] = {}
        # basename -> claude/**/*.py|sh repo-relative paths, for the
        # basename-literal rule below (empty for non-test files / repos with
        # no claude/ tree — see visit_Constant).
        self.claude_basenames = claude_basenames or {}

    def collect_path_consts(self, tree: ast.Module) -> None:
        """First pass: resolve module-level ``NAME = <rooted path expr>``
        assignments to repo-relative path strings (joined segments).
        Directory constants (no file suffix) are kept too — ``HOOKS_DIR /
        hook`` patterns resolve against them."""
        for node in tree.body:
            if not isinstance(node, (ast.Assign, ast.AnnAssign)):
                continue
            targets = node.targets if isinstance(node, ast.Assign) else [node.target]
            if node.value is None:
                continue
            segs = _path_segments(node.value, self.path_consts)
            if segs is None or not segs:
                continue
            joined = "/".join(segs)
            for target in targets:
                if isinstance(target, ast.Name) and target.id.isupper():
                    self.path_consts[target.id] = joined

    def _record(self, kind: str, target: str) -> None:
        edge = {"from": self.rel_path, "kind": kind, "target": target}
        if edge not in self.edges:
            self.edges.append(edge)

    def _div_target(self, expr: ast.expr) -> str | None:
        """Resolve a ``CONST / <something>`` expression. A fully-static
        right side yields a file target; an UNRESOLVABLE right side (loop
        variable, computed name) yields the DIRECTORY as a prefix target —
        the reference is to the whole directory's contents (the
        ``for hook in HOOKS: (HOOKS_DIR / hook).read_text()`` pattern),
        which the resolver matches by prefix."""
        if not isinstance(expr, ast.BinOp) or not isinstance(expr.op, ast.Div):
            return None
        left = _path_segments(expr.left, self.path_consts)
        if left is None or not left:
            return None
        if isinstance(expr.right, ast.Constant) and isinstance(expr.right.value, str):
            joined = "/".join(left + [expr.right.value])
            return joined if joined.endswith(PATH_SUFFIXES) else None
        return "/".join(left) + "/"

    def _argv_paths(self, node: ast.expr) -> list[str]:
        """Repo-relative paths named by an argv list — string literals AND
        resolved module-level path constants (``str(HOOK_PATH)``)."""
        if not isinstance(node, (ast.List, ast.Tuple)):
            return []
        paths: list[str] = []
        for elt in node.elts:
            if isinstance(elt, ast.Constant) and isinstance(elt.value, str):
                if _is_repo_path(elt.value):
                    paths.append(_normalize(elt.value))
            elif isinstance(elt, ast.Name) and elt.id in self.path_consts:
                paths.append(self.path_consts[elt.id])
            elif (
                isinstance(elt, ast.Call)
                and isinstance(elt.func, ast.Name)
                and elt.func.id == "str"
                and len(elt.args) == 1
                and isinstance(elt.args[0], ast.Name)
                and elt.args[0].id in self.path_consts
            ):
                paths.append(self.path_consts[elt.args[0].id])
        return paths

    def visit_Call(self, node: ast.Call) -> None:
        name = _dotted(node.func)
        last = _last(name)

        if last in SUBPROCESS_FUNCS and node.args:
            for arg in self._argv_paths(node.args[0]):
                self._record("invoke", arg)

        if last in BDD_FUNCS:
            for arg in node.args:
                if (
                    isinstance(arg, ast.Constant)
                    and isinstance(arg.value, str)
                    and arg.value.endswith(".feature")
                ):
                    self._record("bdd", _normalize(arg.value))
                elif isinstance(arg, ast.Name) and arg.id in self.path_consts:
                    self._record("bdd", self.path_consts[arg.id])

        if "consult" in last.lower():
            for arg in node.args:
                if (
                    isinstance(arg, ast.Constant)
                    and isinstance(arg.value, str)
                    and _is_repo_path(arg.value)
                ):
                    self._record("consult", _normalize(arg.value))
                elif isinstance(arg, ast.Name) and arg.id in self.path_consts:
                    self._record("consult", self.path_consts[arg.id])

        # Content references — a test that READS a hook/script asserts its
        # content just as surely as one that executes it (2026-08-11:
        # test_session_start_hooks covers inject-codebase-map.sh entirely
        # through read_text assertions on path constants).
        if last in READ_FUNCS and node.args:
            arg0 = node.args[0]
            if isinstance(arg0, ast.Name) and arg0.id in self.path_consts:
                self._record("read", self.path_consts[arg0.id])
            elif (
                isinstance(arg0, ast.Constant)
                and isinstance(arg0.value, str)
                and _is_repo_path(arg0.value)
            ):
                self._record("read", _normalize(arg0.value))
            elif (div := self._div_target(arg0)) is not None:
                self._record("read", div)
        # (HOOKS_DIR / name).read_text() — the file object itself is the
        # attribute value, so the path expression sits under func, not args.
        if last in READ_FUNCS and isinstance(node.func, ast.Attribute):
            if (div := self._div_target(node.func.value)) is not None:
                self._record("read", div)
            elif (
                isinstance(node.func.value, ast.Name)
                and node.func.value.id in self.path_consts
            ):
                self._record("read", self.path_consts[node.func.value.id])

        self.generic_visit(node)

    def visit_Constant(self, node: ast.Constant) -> None:
        """Basename-literal rule: a TEST file UNDER ``tests/claude/`` that
        mentions the exact basename of a ``claude/**/*.py|.sh`` file, as a
        bare string literal, references it — no matter how the path is
        actually assembled at runtime.

        This exists for shapes the call-node walk above cannot see: measured
        2026-08-14, ``tests/claude/test_intercept_guards.py`` binds three
        hooks via ``_HOOKS = Path(__file__).resolve().parents[2] / "claude"
        / "hooks"`` (an anchor the directory-constant machinery does not
        recognize -- ``ROOT_NAMES`` only understands module-level constants
        rooted directly in a name like ``REPO_ROOT``, not
        ``Path(__file__).resolve().parents[N]``) and then calls
        ``str(_HOOKS / hook)`` where ``hook`` is a runtime parameter -- the
        literal basenames (``"intercept_write.py"``, ...) appear only as
        plain string arguments to a local test helper, never as an argv
        element or a path-constant reference.

        Restricted to ``tests/claude/`` referrers specifically (not every
        ``tests/`` module) -- a basename like ``"utils.py"`` or
        ``"__init__.py"`` is common enough that a repo-wide ``tests/``
        gate mints edges from modules that merely mention the string about
        an UNRELATED tree. Measured 2026-08-14:
        ``claude/hooks/error_gate/utils.py`` bound to
        ``test_deepseek_coder_driver.py``, ``test_code_intelligence.py`` and
        others purely because they contain the literal ``"utils.py"`` about
        a *different* utils.py, and with 75 claude/ files sharing the
        ``__init__.py`` basename, every one of those bound to all 75 test
        modules that happen to say ``"__init__.py"`` anywhere -- coincidental
        bindings so dense they SUPPRESSED the ``tests/claude/`` directory
        fallback in ``affected_tests.py`` (a wrong-but-non-empty referrer set
        short-circuited the fallback that would otherwise have fired). Real
        bindings from OUTSIDE ``tests/claude/`` -- e.g. shell-hook tests
        under ``tests/scripts/`` -- are covered instead by the
        ``claude/hooks/*.sh`` manifest rule (``plans/affected-tests-
        manifest.json``), and ``affected_tests.py``'s ``tests/claude/``
        directory-convention fallback is now unconditionally additive for
        every ``claude/`` file, so narrowing this gate cannot silently drop
        coverage the way the false bindings it replaces used to threaten.
        """
        if (
            self.claude_basenames
            and isinstance(node.value, str)
            and self.rel_path.startswith("tests/claude/")
        ):
            for target in self.claude_basenames.get(node.value, ()):
                self._record("basename", target)
        self.generic_visit(node)


def _claude_basename_map(repo_root: Path) -> dict[str, list[str]]:
    """``{basename: [claude/**/*.py|sh repo-relative paths, ...]}``.

    Powers ``visit_Constant`` above. Scoped to the ``claude/`` tree only --
    the one directory whose Python/shell source lives outside
    src/scripts/tests and is invisible to both the import closure and the
    directory-constant reference machinery (2026-08-14 defect: a 3-file
    change to ``claude/hooks/intercept_*.py`` escalated to the full suite
    despite real coverage in ``tests/claude/test_intercept_guards.py``).
    Sorted ``rglob`` keeps output deterministic; a basename that names more
    than one file (``__init__.py`` recurs under several ``claude/`` package
    dirs) maps to every one of them -- ambiguous is safe here, since an
    extra referrer is over-selection, never a missed one.
    """
    claude_dir = repo_root / "claude"
    mapping: dict[str, list[str]] = {}
    if not claude_dir.is_dir():
        return mapping
    for suffix in ("*.py", "*.sh"):
        for p in sorted(claude_dir.rglob(suffix)):
            rel = p.relative_to(repo_root).as_posix()
            mapping.setdefault(p.name, []).append(rel)
    return mapping


def _extract_file(
    args: tuple[Path, str], *, claude_basenames: dict[str, list[str]] | None = None
) -> list[dict[str, str]]:
    """Edges from ONE file — the thread-pool unit of work. Each file parses
    independently (no import graph needed), so this scales linearly on the
    free-threaded build; output is order-independent, sorted by the caller.
    """
    py, rel = args
    try:
        tree = ast.parse(py.read_bytes())
    except SyntaxError, OSError:
        return []
    visitor = _RefEdgeVisitor(rel, claude_basenames=claude_basenames)
    try:
        visitor.collect_path_consts(tree)
        visitor.visit(tree)
    except RecursionError:
        # One pathological file must never take down the whole extraction:
        # that is exactly how a single vendored module silently disabled
        # --affected selection for every run (2026-08-17). Dropping one file's
        # edges costs a little selection precision; raising here costs the
        # feature entirely, and fails the suite with it.
        print(
            f"extract-reference-edges: skipped {rel} (AST nesting too deep)",
            file=sys.stderr,
        )
        return []
    return visitor.edges


# Directory components that are NEVER project source. A vendored dependency
# tree walked as if it were ours is both wrong and unbounded: `.claude/` can
# hold entire agent worktrees, each with its own `.venv`. Measured 2026-08-17 —
# four stale worktrees turned a ~2K-file walk into 76,768 files, and sympy's
# generated `resolvent_lookup.py` nests 569 AST levels. `ast.NodeVisitor` costs
# ~2 Python frames per level (visit -> generic_visit -> visit), so 569 levels
# overflows the default 1000-frame recursion limit and the RecursionError took
# down `--affected` selection AND `tests/scripts/test_affected_manifest_rules.py`
# for every run. Excluding vendored trees fixes the cause; raising the recursion
# limit would not, and is actively unsafe here because `_extract_file` runs in a
# ThreadPoolExecutor on the free-threaded build, where a deeper limit without a
# matching `threading.stack_size` trades a clean RecursionError for a SIGSEGV.
_NON_SOURCE_DIRS = frozenset(
    {
        ".venv",
        "venv",
        "site-packages",
        "node_modules",
        ".git",
        "__pycache__",
        ".mypy_cache",
        ".pytest_cache",
        # Agent worktrees are CHECKOUTS OF THIS REPO AT OTHER COMMITS, parked
        # under `.claude/worktrees/agent-*`. The 2026-08-17 fix above excluded
        # the `.venv` inside them, which cured the 76,768-file walk and the
        # sympy recursion crash, but left their src/, tests/ and scripts/ being
        # indexed as if they were ours.
        #
        # That is a correctness problem, not a size one. Those trees are pinned
        # at OLD commits (measured 2026-08-19: five worktrees at c24da189,
        # 60e7affe x3 and d0fee172), so a test module renamed or deleted since
        # still exists in them — and an edge naming a path that is no longer in
        # the working tree is a phantom, which escalates the whole suite. The
        # same mechanism can also make a change look COVERED by a test only a
        # stale worktree still has. It already distorted the evidence: the
        # affected-set oracle reported "resolved via 5 reference edge(s)" for
        # plans/ast-tests.json where all five cited were worktree copies of one
        # file, so a single real referrer read as five.
        "worktrees",
    }
)


def _is_project_source(path: Path) -> bool:
    """True when *path* is our own source, not a vendored/generated tree."""
    return _NON_SOURCE_DIRS.isdisjoint(path.parts)


def extract_edges(
    source_roots: list[Path], *, repo_root: Path, parallel: bool = True
) -> list[dict[str, str]]:
    """Walk every .py file under ``source_roots`` and return sorted edges."""
    files: list[tuple[Path, str]] = []
    for root in source_roots:
        root = root if root.is_absolute() else repo_root / root
        if not root.exists():
            continue
        for py in sorted(root.rglob("*.py")):
            if not _is_project_source(py):
                continue
            files.append((py, py.relative_to(repo_root).as_posix()))

    worker = functools.partial(
        _extract_file, claude_basenames=_claude_basename_map(repo_root)
    )
    if parallel and len(files) > 1:
        with ThreadPoolExecutor(max_workers=os.cpu_count() or 4) as executor:
            per_file = list(executor.map(worker, files))
    else:
        per_file = [worker(f) for f in files]

    edges = [edge for file_edges in per_file for edge in file_edges]
    edges.sort(key=lambda e: (e["from"], e["kind"], e["target"]))
    return edges


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=__doc__.splitlines()[0] if __doc__ else None
    )
    parser.add_argument(
        "--source",
        nargs="+",
        required=True,
        help="source roots to walk (repo-relative)",
    )
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args(argv)

    repo_root = Path(__file__).resolve().parent.parent
    edges = extract_edges([Path(s) for s in args.source], repo_root=repo_root)
    # Atomic: affected_tests.py regenerates this mid-run while the tree is
    # being staged; a truncate-in-place write keeps the inode and SIGBUSes a
    # reader that already mmap'd it. Byte-identical output (indent=1 plus the
    # trailing newline the previous writer appended).
    from orchestrata_ultimate.core.atomic_json import dump_json_atomic

    dump_json_atomic(
        args.output,
        {"version": 1, "edges": edges},
        indent=1,
        trailing_newline=True,
    )
    print(f"[ref-edges] {len(edges)} edges -> {args.output}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())