scripts.add_traces_from_evidence

Reference › scripts › add_traces_from_evidence

A command-line tool that writes @traces decorators onto source functions and classes, using test citations as evidence. It reads the requirement definitions under the specs directory and harvests the requirement ids that tests cite. It then joins each citing test to the code it calls, through the BodyMatcher from body_matcher, the JSONASTLoader from json_loader and the InheritanceMap from inheritance. No other module or test imports it; a human runs it.

It is invoked as a plain script with the options --specs (default specs), --tests (default tests), --src-ast (default plans/ast-ultimate.json), --tests-ast (default plans/ast-tests.json) and --apply. Without --apply it reports what it would change; with it, it edits the source files in place.

main first maps each test function to the requirement ids cited inside its span (enclosing_test_functions). Citations of ids that no spec defines are skipped, so an undefined id is never written. For each citing test, the calls found in its own body are resolved to definitions only when the test's imports make the target unambiguous: a Class.method call resolves to the method (ancestors count), a bare imported function resolves to itself, and a constructor call counts as nothing. Targets must live under the main source package. defining_spec_number finds the spec that defines a requirement, and process_file adds or merges the decorators in one file and returns three counters.

At a glance

Class diagram

No classes; see Functions.

Dependency diagram

API

Functions

Source

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

"""Write @traces from EVIDENCE: test citations joined to code via TestMatcher.

The only trustworthy retroactive trace link is one derived from evidence a
human already authored (AGENTS.md "@traces Conventions"; strategy review
2026-07-09). This tool builds exactly that chain, deterministically:

  requirement R  --cited in-->  test function T   (authored claim: the REQ id
                                                   appears inside T's body or
                                                   docstring)
  T  --direct body call-->  source def M          (BodyMatcher.direct_call_index:
                                                   T's own body calls M, resolved
                                                   through imports/var types;
                                                   fixture calls excluded)
  M  gets  @traces(spec=<R's defining spec>, reqs=[R])

Discipline (mirrors trace-lint's rules so output is lint-clean by
construction):
- Only citations INSIDE a test function/class body or docstring count;
  module-level mentions are too coarse for def-level attribution.
- R must be defined in some spec's requirements.md (GHOST ids are skipped).
- spec= is R's defining spec (prefix family preferred) — membership-true.
- story= is the US-segment embedded in R's own id when present, else "".
- Existing decorator already carrying R -> skipped (idempotent).
- Same-spec existing @traces on M -> R is appended to its reqs; otherwise a
  stacked decorator line is inserted (mixed provenance never shares one
  decorator).
- Dry-run by default; --apply writes.

Usage:
  python scripts/add_traces_from_evidence.py               # report only
  python scripts/add_traces_from_evidence.py --apply
"""

from __future__ import annotations

import argparse
import ast
import re
import sys
from collections import defaultdict
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

from generate_rtm_evidence import (  # noqa: E402
    harvest_test_citations,
    load_requirements,
)

# Canonical story segment only (RUL-031): ``US``, a hyphen, three digits. The
# capture is emitted VERBATIM as the ``story=`` value, because this script
# writes @traces decorators and every story label in the corpus is the
# zero-padded ``US-NNN`` form. The previous pattern captured the legacy
# ``US\d+`` and rebuilt the label by string surgery (``"US-" + g[2:]``), which
# produced un-padded labels like ``US-7`` and matched nothing at all once the
# corpus was migrated.
_STORY_IN_ID_RE = re.compile(r"REQ-\d+[a-z]?-(US-\d{3})-", re.IGNORECASE)

# Two-lane ruling (AGENTS.md "@traces Conventions"): infrastructure code is
# deliberately untraced unless a spec names an explicit NFR for it. Tests
# call these as SETUP (build a client, spin a pool), which is exercise, not
# implementation — never attribute functional reqs here.
_INFRA_LANE = (
    "src/orchestrata_ultimate/core/base/",
    "src/orchestrata_ultimate/core/redis_factory.py",
    "src/orchestrata_ultimate/core/graph_factory.py",
    "src/orchestrata_ultimate/core/threading/",
    "src/orchestrata_ultimate/core/lru_cache.py",
    "src/orchestrata_ultimate/core/logging.py",
    "src/orchestrata_ultimate/core/paths.py",
)


def _story_for(rid: str) -> str:
    m = _STORY_IN_ID_RE.match(rid)
    return m.group(1) if m else ""


def _spec_number(spec_dir_name: str) -> str:
    m = re.match(r"^(\d+[a-z]?)", spec_dir_name)
    return m.group(1) if m else spec_dir_name


def defining_spec_number(rid: str, reqs: dict[str, dict]) -> str | None:
    entry = reqs.get(rid)
    if not entry:
        return None
    fam = re.match(r"REQ-(\d+[a-z]?)-", rid)
    if fam:
        for spec_name in entry["specs"]:
            if _spec_number(spec_name).lstrip("0") == fam.group(1).lstrip("0"):
                return _spec_number(spec_name)
    return _spec_number(entry["specs"][0])


def enclosing_test_functions(
    test_file: Path, lines_cited: dict[int, set[str]]
) -> dict[tuple[str, str | None], set[str]]:
    """(test_name, class_name) -> req ids, for citations inside that test's
    span (body, docstring, decorators). Class-docstring citations attach to
    every test in the class."""
    try:
        tree = ast.parse(test_file.read_text(encoding="utf-8"))
    except OSError, SyntaxError, UnicodeDecodeError:
        return {}
    out: dict[tuple[str, str | None], set[str]] = defaultdict(set)

    def span_reqs(node: ast.stmt | ast.expr) -> set[str]:
        start = min(
            [node.lineno] + [d.lineno for d in getattr(node, "decorator_list", [])]
        )
        got: set[str] = set()
        for ln, rids in lines_cited.items():
            if start <= ln <= (node.end_lineno or node.lineno):
                got |= rids
        return got

    for node in tree.body:
        if isinstance(node, ast.ClassDef):
            class_doc_reqs: set[str] = set()
            doc = ast.get_docstring(node)
            if doc and node.body and isinstance(node.body[0], ast.Expr):
                class_doc_reqs = span_reqs(node.body[0])
            for item in node.body:
                if isinstance(
                    item, (ast.FunctionDef, ast.AsyncFunctionDef)
                ) and item.name.startswith("test_"):
                    rids = span_reqs(item) | class_doc_reqs
                    if rids:
                        out[(item.name, node.name)] |= rids
        elif isinstance(
            node, (ast.FunctionDef, ast.AsyncFunctionDef)
        ) and node.name.startswith("test_"):
            rids = span_reqs(node)
            if rids:
                out[(node.name, None)] |= rids
    return out


def process_file(
    p: Path,
    wants: list[tuple[int, str, dict]],
    reqs: dict[str, dict],
    apply: bool,
) -> tuple[int, int, int]:
    """Add/merge @traces decorators in one file for the wanted defs.

    ``wants`` is [(line, def_name, {"reqs": set, "evidence": set})]. Returns
    (added, merged, skipped). Writes only when ``apply`` and something
    changed; the result is always re-parseable (ast-anchored edits).
    """
    src = p.read_text(encoding="utf-8")
    tree = ast.parse(src)
    lines = src.splitlines(keepends=True)
    edits: list[tuple[int, str]] = []  # (insert_at_line0, text); merges edit `lines`
    added = merged = skipped = 0
    # any existing binding of `traces` counts (core.base.traces,
    # core.traceability, aliases) — a second import would shadow/F811
    needs_import = not any(
        isinstance(n, ast.ImportFrom)
        and any((a.asname or a.name) == "traces" for a in n.names)
        for n in tree.body
    )
    for node in ast.walk(tree):
        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            continue
        for want_line, want_name, data in wants:
            if node.name != want_name or abs(node.lineno - want_line) > 5:
                continue
            existing_reqs: set[str] = set()
            same_spec_dec: dict[str, ast.Call] = {}
            for dec in node.decorator_list:
                if isinstance(dec, ast.Call) and (
                    (isinstance(dec.func, ast.Name) and dec.func.id == "traces")
                    or (
                        isinstance(dec.func, ast.Attribute)
                        and dec.func.attr == "traces"
                    )
                ):
                    kw = {k.arg: k.value for k in dec.keywords}
                    rv = kw.get("reqs")
                    if isinstance(rv, (ast.List, ast.Tuple)):
                        existing_reqs |= {
                            e.value
                            for e in rv.elts
                            if isinstance(e, ast.Constant) and isinstance(e.value, str)
                        }
                    sv = kw.get("spec")
                    if isinstance(sv, ast.Constant):
                        same_spec_dec[str(sv.value)] = dec
            new_by_spec: dict[str, list[str]] = defaultdict(list)
            for rid in sorted(data["reqs"] - existing_reqs):
                spec = defining_spec_number(rid, reqs)
                if spec:
                    new_by_spec[spec].append(rid)
            if not new_by_spec:
                skipped += 1
                continue
            indent_m = re.match(r"\s*", lines[node.lineno - 1])
            indent = indent_m.group(0) if indent_m else ""
            dec_start = min([node.lineno] + [d.lineno for d in node.decorator_list])
            for spec, rids in sorted(new_by_spec.items()):
                sdec = same_spec_dec.get(spec)
                rv_node = (
                    {k.arg: k.value for k in sdec.keywords}.get("reqs")
                    if sdec is not None
                    else None
                )
                if isinstance(rv_node, (ast.List, ast.Tuple)):
                    # merge into the existing same-spec decorator's reqs
                    # list — anchored AFTER the last element (never before
                    # the closing bracket: a multi-line list's trailing
                    # comma would produce ", ," there)
                    body = ", ".join(f'"{r}"' for r in rids)
                    if rv_node.elts:
                        anchor = rv_node.elts[-1]
                        end_ln = (anchor.end_lineno or anchor.lineno) - 1
                        end_col = anchor.end_col_offset
                        insert = ", " + body
                    else:
                        end_ln = (rv_node.end_lineno or rv_node.lineno) - 1
                        end_col = (rv_node.end_col_offset or 1) - 1
                        insert = body
                    ln = lines[end_ln]
                    lines[end_ln] = ln[:end_col] + insert + ln[end_col:]
                    merged += 1
                else:
                    story = next((s for s in (_story_for(r) for r in rids) if s), "")
                    st = f' story="{story}",' if story else ""
                    text = (
                        f'{indent}@traces(spec="{spec}",{st} '
                        f"reqs=[{', '.join(f'\"{r}\"' for r in rids)}])\n"
                    )
                    edits.append((dec_start - 1, text))
                    added += 1
            for spec, rids in sorted(new_by_spec.items()):
                ev = sorted(data["evidence"])[0] if data["evidence"] else "?"
                print(
                    f"{p}:{node.lineno} {node.name} += "
                    f"[{', '.join(rids)}] (spec {spec}; {ev})"
                )
    if apply and (edits or merged):
        for at, text in sorted(edits, reverse=True):
            lines.insert(at, text)
        if needs_import:
            out = "".join(lines)
            tree2 = ast.parse(out)
            last_import = max(
                (
                    n.end_lineno or n.lineno
                    for n in tree2.body
                    if isinstance(n, (ast.Import, ast.ImportFrom))
                ),
                default=0,
            )
            ls = out.splitlines(keepends=True)
            ls.insert(
                last_import,
                "from orchestrata_ultimate.core.base.traces import traces\n",
            )
            lines = ls
        p.write_text("".join(lines))
    return added, merged, skipped


def main() -> int:
    ap = argparse.ArgumentParser(description=(__doc__ or "").splitlines()[0])
    ap.add_argument("--specs", default="specs")
    ap.add_argument("--tests", default="tests")
    ap.add_argument("--src-ast", default="plans/ast-ultimate.json")
    ap.add_argument("--tests-ast", default="plans/ast-tests.json")
    ap.add_argument("--apply", action="store_true")
    args = ap.parse_args()

    from orchestrata_ast.inheritance import InheritanceMap
    from orchestrata_tools.parser.json_loader import JSONASTLoader
    from orchestrata_tools.test_analyzer.body_matcher import BodyMatcher

    reqs = load_requirements(Path(args.specs))
    cites = harvest_test_citations(Path(args.tests))

    # citations per test FILE: line -> req ids
    lines_by_file: dict[str, dict[int, set[str]]] = defaultdict(
        lambda: defaultdict(set)
    )
    for rid, cs in cites.items():
        if rid not in reqs:
            continue  # GHOST — never write an undefined id
        for c in cs:
            lines_by_file[c["file"]][c["line"]].add(rid)

    # test function -> req ids (citation INSIDE the function/class span).
    # Keyed by the AST snapshot's dotted module path (tests/a/b.py -> a.b)
    # so it joins against TestEntry.module.
    test_reqs: dict[tuple[str, str, str | None], set[str]] = {}
    for tf, line_map in lines_by_file.items():
        dotted = str(Path(tf).relative_to(args.tests).with_suffix("")).replace("/", ".")
        for (tname, tcls), rids in enclosing_test_functions(Path(tf), line_map).items():
            test_reqs[(dotted, tname, tcls)] = rids

    # direct-call evidence: each citing test's OWN body calls, resolved
    test_ast = JSONASTLoader().load(args.tests_ast)
    src_ast = JSONASTLoader().load(args.src_ast)
    bm = BodyMatcher(test_ast, src_ast)
    call_index = bm.direct_call_index()

    # source lookup tables
    functions_by_name: dict[str, list] = defaultdict(list)
    for fn in src_ast.all_functions:
        functions_by_name[fn.name].append(fn)
    classes_by_name: dict[str, list] = defaultdict(list)
    for cls in src_ast.all_classes:
        classes_by_name[cls.name].append(cls)

    # per test module: imported name -> defining module (from-imports only)
    imports_map: dict[str, dict[str, str]] = defaultdict(dict)
    for mod in test_ast.modules:
        for imp in mod.imports:
            if imp.is_from and imp.module.startswith("orchestrata_ultimate"):
                for name in imp.names:
                    imports_map[mod.path][name] = imp.module

    def _mod_matches(candidate_mod: str, imported_mod: str) -> bool:
        im = imported_mod.removeprefix("orchestrata_ultimate").lstrip(".")
        if not im:
            return True
        return (
            candidate_mod == im
            or candidate_mod.startswith(im + ".")
            or im.startswith(candidate_mod + ".")
        )

    def _resolve_defs(test_mod: str, name: str) -> list:
        """Same-name defs disambiguated by the test's own import; with no
        import record, only a codebase-unique name counts as evidence."""
        cands = classes_by_name.get(name, []) + functions_by_name.get(name, [])
        im = imports_map.get(test_mod, {}).get(name)
        if im is not None:
            return [c for c in cands if _mod_matches(c.module_path, im)]
        return cands if len(cands) == 1 else []

    inheritance = InheritanceMap.from_codebase(src_ast)

    def _ancestor_methods(cls_name: str, meth_name: str) -> list:
        """The nearest ancestor (declaration-order DFS) defining the method."""
        for anc in inheritance.iter_mro(cls_name):
            for bc in classes_by_name.get(anc, []):
                ms = [m for m in bc.methods if m.name == meth_name]
                if ms:
                    return ms
        return []

    def resolve_call(test_mod: str, call: str, imported: set[str]) -> list:
        """A call is evidence only when it resolves unambiguously through the
        test's own imports. Class.method resolves to the method (ancestors
        count); a bare imported function resolves to itself. Constructor
        calls resolve to NOTHING — building an object as test input is not
        evidence the class implements the requirement."""
        parts = call.split(".")
        if len(parts) >= 2:
            cls_name, meth_name = parts[0], parts[-1]
            if cls_name not in imported:
                return []
            out = []
            for target in _resolve_defs(test_mod, cls_name):
                if not hasattr(target, "methods"):
                    continue  # a function, not a class
                ms = [m for m in target.methods if m.name == meth_name]
                out += ms or _ancestor_methods(target.name, meth_name)
            return out
        name = parts[0]
        if name not in imported:
            return []
        return [t for t in _resolve_defs(test_mod, name) if not hasattr(t, "methods")]

    additions: dict[tuple[str, str, int], dict] = {}
    for key, rids in test_reqs.items():
        calls, imported = call_index.get(key, (set(), set()))
        for call in sorted(calls):
            for target in resolve_call(key[0], call, imported):
                if not target.file_path.startswith("src/orchestrata_ultimate"):
                    continue
                if target.file_path.startswith(_INFRA_LANE):
                    continue
                entry = additions.setdefault(
                    (target.file_path, target.name, target.line),
                    {"reqs": set(), "evidence": set()},
                )
                entry["reqs"] |= rids
                entry["evidence"].add(f"{key[0]}::{key[1]} calls {call}")

    if not additions:
        print("no evidence-backed additions found")
        return 0

    # Apply per file, ast-precise, bottom-up.
    by_file: dict[str, list] = defaultdict(list)
    for (fpath, name, line), data in additions.items():
        by_file[fpath].append((line, name, data))

    added, merged, skipped = 0, 0, 0
    for fpath in sorted(by_file):
        p = Path(fpath)
        if not p.exists():
            continue
        a, m, s = process_file(p, by_file[fpath], reqs, apply=args.apply)
        added, merged, skipped = added + a, merged + m, skipped + s

    print(
        f"\n{'APPLIED' if args.apply else 'DRY-RUN'}: "
        f"{added} new decorators, {merged} same-spec merges, "
        f"{skipped} already-present, across {len(by_file)} files"
    )
    return 0


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