scripts.apply_docstring_traces

Reference › scripts › apply_docstring_traces

Converts requirement mentions that already sit in test docstrings and comments into @traces decorators. It is a one-off migration helper run by a human. The only internal module it relates to is traces, which defines the decorator it emits, and tests/scripts/test_apply_docstring_traces.py imports it.

Invocation: two positional arguments, a test package directory such as tests/kernel and a spec number such as 001, plus --dry-run to print changes without writing. It exits with an error if the directory does not exist.

main walks every test_*.py file under the package and calls _process_file on each. For each function or class, _extract_req_ids_from_docstring finds the requirement ids mentioned, _has_traces_decorator skips nodes that already carry the decorator, _infer_story derives a story from the ids, and _format_traces_decorator produces the decorator text. Per-file errors go to stderr; the final line reports the number of decorators added and files touched.

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/apply_docstring_traces.py

"""Convert existing docstring/comment REQ mentions into @traces decorators.

This script is conservative: it ONLY adds @traces to test functions/classes
that already mention requirement IDs in their docstrings or inline comments.
It does not guess. It verifies syntax after every edit.
"""

from __future__ import annotations

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

# The SAME validator ``@traces`` enforces at import time. Importing it rather
# than restating the grammar is the point: a second copy can drift, and a
# generator whose idea of "valid" drifts from the decorator's writes code that
# raises on collection. There is exactly one grammar, and this reads it.
from orchestrata_ultimate.core.base.traces import _validate_req_id

_REQ_MENTION_RE = re.compile(r"REQ-[0-9][A-Za-z0-9-]*")
_TRACES_IMPORT_LINE = "from orchestrata_ultimate.core.base.traces import traces\n"

#: A module that declares its REQ ids are fixture data is not making coverage
#: claims, and must not have claims written into it. tests/scripts/
#: test_generate_rtm_evidence_ts.py carries this marker and was scraped anyway.
_FIXTURE_IDS_MARKER = "rtm-fixture-ids"


def _extract_req_ids_from_docstring(text: str) -> list[str]:
    """Return sorted unique requirement IDs found in *text* (docstring/comment).

    VERIFY BEFORE WRITING. A mention becomes a citation only if it is a real
    requirement id under the grammar ``@traces`` itself enforces, so this
    generator cannot emit a decorator that raises when the module is imported.

    The regex alone is not that test. Prose names a SPEC by a bare ``REQ-NNN``
    -- "REQ-022 measured 39 of 39 UNLINKED" means SPEC-022's 39 requirements,
    not a requirement called REQ-022 -- and the unvalidated match wrote
    ``@traces(spec="000", reqs=["REQ-022"])`` into a test module. ``@traces``
    raises ValueError on that id at DECORATION time, i.e. at import, so pytest
    could not even collect the module and every batch containing it failed
    (measured 2026-08-26: eight reported batch failures, one cause).
    """
    ids = sorted({m.group(0) for m in _REQ_MENTION_RE.finditer(text)})
    # Filter out fixture sentinel ids, then the grammar itself.
    return [rid for rid in ids if rid != "REQ-000-000" and _validate_req_id(rid)]


def _has_traces_decorator(
    node: ast.FunctionDef | ast.ClassDef | ast.AsyncFunctionDef,
) -> bool:
    for dec in node.decorator_list:
        if isinstance(dec, ast.Name) and dec.id == "traces":
            return True
        if isinstance(dec, ast.Call):
            func = dec.func
            if isinstance(func, ast.Name) and func.id == "traces":
                return True
            if isinstance(func, ast.Attribute) and func.attr == "traces":
                return True
    return False


def _format_traces_decorator(spec: str, story: str | None, reqs: list[str]) -> str:
    if story:
        return f'@traces(spec="{spec}", story="{story}", reqs={reqs})\n'
    return f'@traces(spec="{spec}", reqs={reqs})\n'


def _infer_story(req_ids: list[str]) -> str | None:
    """Infer story from REQ id segments if possible, e.g. REQ-001-US-002-015 -> US-002."""
    stories: set[str] = set()
    for rid in req_ids:
        parts = rid.split("-")
        # REQ-NNN-US-NNN-NNN or REQ-NNN-FRNN-NNN
        if len(parts) >= 4 and parts[2] in ("US", "FR"):
            stories.add(f"{parts[2]}-{parts[3]}")
    if len(stories) == 1:
        return stories.pop()
    return None


def _process_file(
    path: Path, spec: str, dry_run: bool = False
) -> tuple[bool, int, list[str]]:
    """Return (changed, decorators_added, errors)."""
    try:
        source = path.read_text(encoding="utf-8")
    except OSError as e:
        return False, 0, [f"read error: {e}"]

    if _FIXTURE_IDS_MARKER in source:
        # The module states that its REQ ids are fixture/sample data for
        # exercising traceability machinery. Scraping them writes a coverage
        # claim into the one place that has declared it makes none.
        return False, 0, []

    try:
        tree = ast.parse(source)
    except SyntaxError as e:
        return False, 0, [f"pre-existing syntax error: {e}"]

    lines = source.splitlines(keepends=True)
    # Ensure each line ends with newline for stable insertion logic
    if lines and not lines[-1].endswith("\n"):
        lines[-1] += "\n"

    # Collect targets: (insert_line_index, indent, story, reqs, node_name)
    targets: list[tuple[int, str, str | None, list[str], str]] = []

    for node in ast.walk(tree):
        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            continue
        if not node.name.startswith("test") and not (
            isinstance(node, ast.ClassDef) and node.name.startswith("Test")
        ):
            continue
        if _has_traces_decorator(node):
            continue

        docstring = ast.get_docstring(node)
        req_ids: list[str] = []
        if docstring:
            req_ids = _extract_req_ids_from_docstring(docstring)

        if not req_ids:
            continue

        story = _infer_story(req_ids)
        # Insert BEFORE decorator list. If no decorators, insert at function def line.
        insert_lineno = node.lineno
        # Find the first decorator line if any
        for dec in node.decorator_list:
            if hasattr(dec, "lineno") and dec.lineno < insert_lineno:
                insert_lineno = dec.lineno

        # Compute indentation of the node definition line
        def_line = lines[insert_lineno - 1]
        indent = def_line[: len(def_line) - len(def_line.lstrip())]
        targets.append((insert_lineno - 1, indent, story, req_ids, node.name))

    if not targets:
        return False, 0, []

    # Sort targets by line descending so insertion indices remain valid
    targets.sort(key=lambda t: t[0], reverse=True)

    new_lines = list(lines)
    added = 0
    for insert_idx, indent, story, req_ids, name in targets:
        dec_line = indent + _format_traces_decorator(spec, story, req_ids)
        new_lines.insert(insert_idx, dec_line)
        added += 1

    # Add import if missing (check via AST, not string search, so fixture
    # strings containing the import text are not mistaken for real imports).
    has_import = False
    for top in tree.body:
        if isinstance(top, ast.ImportFrom):
            module = top.module or ""
            if (
                module.endswith(".traces")
                or module == "orchestrata_ultimate.core.base.traces"
            ):
                for alias in top.names:
                    if alias.name == "traces":
                        has_import = True
                        break
        if has_import:
            break
    if not has_import and added:
        # Find the last line of any top-level import using AST end_lineno
        # (handles multi-line `from x import (...)` correctly).
        last_import_end = 0
        for top in tree.body:
            if isinstance(top, (ast.Import, ast.ImportFrom)):
                end = getattr(top, "end_lineno", top.lineno)
                if end and end > last_import_end:
                    last_import_end = end
        # Convert 1-based end_lineno to 0-based insert index in new_lines
        insert_import_idx = last_import_end
        new_lines.insert(insert_import_idx, _TRACES_IMPORT_LINE)

    new_source = "".join(new_lines)

    # Verify syntax
    try:
        ast.parse(new_source)
    except SyntaxError as e:
        return False, 0, [f"generated syntax error in {path}: {e}"]

    if dry_run:
        return False, added, []

    path.write_text(new_source, encoding="utf-8")
    return True, added, []


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Add @traces decorators from existing docstring REQ mentions."
    )
    parser.add_argument("package", help="Test package directory, e.g. tests/kernel")
    parser.add_argument("spec", help="Spec number, e.g. 001")
    parser.add_argument(
        "--dry-run", action="store_true", help="Print changes without writing"
    )
    args = parser.parse_args()

    pkg = Path(args.package)
    if not pkg.is_dir():
        print(f"ERROR: {pkg} is not a directory", file=sys.stderr)
        return 1

    total_added = 0
    total_files = 0
    for py in sorted(pkg.rglob("test_*.py")):
        changed, added, errors = _process_file(py, args.spec, dry_run=args.dry_run)
        if errors:
            for err in errors:
                print(f"  ERROR {py}: {err}", file=sys.stderr)
        if added:
            total_files += 1
            total_added += added
            print(f"{'[dry-run] ' if args.dry_run else ''}{py}: +{added} decorators")

    print(f"\nTotal: {total_added} decorators across {total_files} files")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())