orchestrata_ultimate.core.tracing.scanner

Reference › orchestrata_ultimate › core › tracing › scanner

The AST scanner that extracts @traces annotations from source trees without importing the code, used by the session graph-sync stage and by scripts that build the requirements traceability matrix, check for feature loss and sync traces to the graph. The TraceInfo records it returns are defined in tracing.traceability.

The design separates parsing from scanning so that a tree is read and parsed once. ModuleCorpus.load(root, module_filter, tracer) reads every wanted module under a directory exactly once and holds the parsed ParsedModule records (path, module name, text, syntax tree) in sorted-path order; it raises ValueError if the root is not a directory, and records the modules a filter rejected in skipped. A ModuleFilter decides which modules to parse: AcceptAllFilter takes all, and RequirementMentionFilter takes only modules whose text contains REQ-, which was measured to cut most of the parsing time on the large tests tree. TraceScanner then walks the trees: collect_from(corpus) scans already-parsed modules, and collect_traces(root_dir) streams a directory without keeping a corpus. For classes, functions and methods, _extract_traces_from_decorator reads the keyword arguments of the decorator call, and a method inherits the class's annotation as a parent trace. The scanner takes an injected Tracer for spans, defaulting to the process-wide one.

At a glance

Class diagram

Dependency diagram

API

Classes

ParsedModule (dataclass)

Decorators: dataclass(frozen=True) · Line 198

One module read and parsed exactly once (WI-FLOSS, tasks 578/581).

ModuleFilter (protocol)

Bases: Protocol · Line 207

Strategy deciding, from the raw text, whether a module is worth parsing.

AcceptAllFilter

Line 215

Default strategy: every readable module is parsed (legacy behaviour).

RequirementMentionFilter

Line 223

Accept only modules whose text contains the literal REQ-.

ModuleCorpus

Line 269

The parsed modules under one root, built once and shared by consumers.

TraceScanner

Line 327

Collects @traces annotations from a corpus or a directory.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/core/tracing/scanner.py

"""AST scanner for @traces decorator annotations.

Walks Python source trees and extracts TraceInfo metadata from
@traces decorators on classes and functions/methods.
"""

import ast
import logging
from collections.abc import Iterable, Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol

from orchestrata_ultimate.core.tracing.traceability import TraceInfo
from orchestrata_ultimate.core.tracing.tracer import Tracer

logger = logging.getLogger(__name__)


def _extract_string_value(node: ast.expr) -> str:
    """Extract a string value from an AST node."""
    if isinstance(node, ast.Constant) and isinstance(node.value, str):
        return node.value
    return ""


def _extract_list_values(node: ast.expr) -> list[str]:
    """Extract a list of strings from an AST List node."""
    if isinstance(node, ast.List):
        return [
            _extract_string_value(elt)
            for elt in node.elts
            if isinstance(elt, ast.Constant) and isinstance(elt.value, str)
        ]
    return []


def _extract_traces_from_decorator(
    call_node: ast.Call,
) -> dict[str, Any] | None:
    """Extract traces kwargs from a decorator Call node."""
    func = call_node.func
    if isinstance(func, ast.Name) and func.id == "traces":
        pass
    elif isinstance(func, ast.Attribute) and func.attr == "traces":
        pass
    else:
        return None

    result: dict[str, Any] = {"spec": "", "story": "", "reqs": []}

    for kw in call_node.keywords:
        if kw.arg == "spec":
            result["spec"] = _extract_string_value(kw.value)
        elif kw.arg == "story":
            result["story"] = _extract_string_value(kw.value)
        elif kw.arg == "reqs":
            result["reqs"] = _extract_list_values(kw.value)

    return result if result["reqs"] else None


def _file_path_to_module(file_path: Path, root_dir: Path) -> str:
    """Convert a file path to a dotted module name relative to root."""
    try:
        relative = file_path.relative_to(root_dir)
    except ValueError:
        relative = file_path
    parts = list(relative.with_suffix("").parts)
    return ".".join(parts)


def _is_method(
    func_node: ast.FunctionDef | ast.AsyncFunctionDef,
    tree: ast.AST,
) -> bool:
    """Check if a function node is a method (inside a class)."""
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            for item in node.body:
                if item is func_node:
                    return True
    return False


def _scan_function(
    func_node: ast.FunctionDef | ast.AsyncFunctionDef,
    module: str,
    class_name: str,
    file_path: str,
    parent_trace: dict[str, Any] | None = None,
) -> list[TraceInfo]:
    """Scan a function/method for @traces decorators."""
    results: list[TraceInfo] = []

    for dec in func_node.decorator_list:
        if isinstance(dec, ast.Call):
            data = _extract_traces_from_decorator(dec)
            if data:
                spec = data["spec"] or (parent_trace["spec"] if parent_trace else "")
                story = data["story"] or (parent_trace["story"] if parent_trace else "")
                results.extend(
                    TraceInfo(
                        spec=spec,
                        story=story,
                        reqs=[req_id],
                        module=module,
                        class_name=class_name,
                        method=func_node.name,
                        file_path=file_path,
                        line_number=func_node.lineno,
                    )
                    for req_id in data["reqs"]
                )

    return results


def _scan_class(
    class_node: ast.ClassDef,
    module: str,
    file_path: str,
) -> list[TraceInfo]:
    """Scan a class definition for @traces decorators."""
    results: list[TraceInfo] = []
    class_trace_data = None

    for dec in class_node.decorator_list:
        if isinstance(dec, ast.Call):
            data = _extract_traces_from_decorator(dec)
            if data:
                class_trace_data = data
                results.extend(
                    TraceInfo(
                        spec=data["spec"],
                        story=data["story"],
                        reqs=[req_id],
                        module=module,
                        class_name=class_node.name,
                        method="",
                        file_path=file_path,
                        line_number=class_node.lineno,
                    )
                    for req_id in data["reqs"]
                )

    for item in class_node.body:
        if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
            results.extend(
                _scan_function(
                    item,
                    module,
                    class_node.name,
                    file_path,
                    class_trace_data,
                )
            )

    return results


def _scan_module_nodes(
    tree: ast.Module,
    module_name: str,
    file_path: str,
) -> list[TraceInfo]:
    """Scan an AST module for @traces on classes and top-level functions.

    Args:
        tree: Parsed AST module.
        module_name: Dotted module name for trace metadata.
        file_path: Absolute file path for trace metadata.

    Returns:
        List of TraceInfo objects found in this module.
    """
    # The direct-body items of every class, computed ONCE. ``_is_method`` walks
    # the whole tree per function, which made this scan O(functions x nodes):
    # measured 2026-10-08, 272s of a 316s ``harvest_test_citations`` profile
    # over ~2,060 test modules (9,093 ``_is_method`` calls); same answer here.
    method_ids = {
        id(item)
        for node in ast.walk(tree)
        if isinstance(node, ast.ClassDef)
        for item in node.body
    }
    results: list[TraceInfo] = []
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            results.extend(_scan_class(node, module_name, file_path))
        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            if id(node) not in method_ids:
                results.extend(_scan_function(node, module_name, "", file_path))
    return results


@dataclass(frozen=True)
class ParsedModule:
    """One module read and parsed exactly once (WI-FLOSS, tasks 578/581)."""

    path: Path
    module_name: str
    text: str
    tree: ast.Module


class ModuleFilter(Protocol):
    """Strategy deciding, from the raw text, whether a module is worth parsing."""

    def wants(self, path: Path, text: str) -> bool:
        """True when *path* (whose source is *text*) belongs in the corpus."""
        ...


class AcceptAllFilter:
    """Default strategy: every readable module is parsed (legacy behaviour)."""

    def wants(self, path: Path, text: str) -> bool:
        """Accept every module."""
        return True


class RequirementMentionFilter:
    """Accept only modules whose text contains the literal ``REQ-``.

    Every citation form the RTM harvesters understand (docstring/comment
    mention and ``@traces(reqs=[...])`` literal) contains ``REQ-``, so a
    module without it can yield nothing. Measured 2026-10-08 over the
    ~2,060-module tests tree: most modules mention no requirement, and
    parsing them dominated ``check_feature_loss.run_check`` (harvest ~63s
    solo after the first pre-filter, with a second full parse in
    ``collect_traces``).
    """

    def wants(self, path: Path, text: str) -> bool:
        """Accept a module only when it mentions a requirement id."""
        return "REQ-" in text


def _stream_modules(
    root: Path, module_filter: ModuleFilter, skipped: list[Path] | None = None
) -> Iterator[ParsedModule]:
    """Yield each module under *root* that is readable, wanted and parses.

    ``sorted(rglob)`` order, ``__pycache__`` excluded, unreadable/unparsable
    modules logged and skipped — exactly the walk ``collect_traces`` always did.
    """
    for py_file in sorted(root.rglob("*.py")):
        if "__pycache__" in str(py_file):
            continue
        try:
            text = py_file.read_text(encoding="utf-8")
            if not module_filter.wants(py_file, text):
                if skipped is not None:
                    skipped.append(py_file)
                continue
            tree = ast.parse(text, filename=str(py_file))
        except (SyntaxError, UnicodeDecodeError) as exc:
            logger.warning("Skipping %s: %s", py_file, exc)
            continue
        yield ParsedModule(
            path=py_file,
            module_name=_file_path_to_module(py_file, root.parent),
            text=text,
            tree=tree,
        )


class ModuleCorpus:
    """The parsed modules under one root, built once and shared by consumers.

    ``load`` does one ``read_text`` and (for wanted modules) one ``ast.parse``
    per file; every consumer then iterates ``modules()`` instead of
    re-reading the tree. Composes an injected :class:`Tracer`.
    """

    def __init__(
        self,
        root: Path,
        modules: list[ParsedModule],
        skipped: list[Path],
        tracer: Tracer,
    ) -> None:
        """Hold the already-parsed *modules* of *root* (use :meth:`load`)."""
        self._root = root
        self._modules = modules
        self._skipped = tuple(skipped)
        self._tracer = tracer

    @classmethod
    def load(
        cls,
        root: Path,
        module_filter: ModuleFilter | None = None,
        tracer: Tracer | None = None,
    ) -> ModuleCorpus:
        """Read and parse every wanted module under *root* exactly once."""
        root = Path(root)
        if not root.is_dir():
            raise ValueError(f"Not a directory: {root}")
        active_tracer = tracer if tracer is not None else Tracer.default()
        skipped: list[Path] = []
        with active_tracer.span("ModuleCorpus.load"):
            modules = list(
                _stream_modules(root, module_filter or AcceptAllFilter(), skipped)
            )
        return cls(root, modules, skipped, active_tracer)

    @property
    def root(self) -> Path:
        """The directory this corpus was loaded from."""
        return self._root

    def modules(self) -> Iterator[ParsedModule]:
        """Iterate the parsed modules (stable, sorted-path order)."""
        return iter(self._modules)

    def __iter__(self) -> Iterator[ParsedModule]:
        """Iterate the parsed modules; lets a corpus feed any consumer."""
        return self.modules()

    def skipped(self) -> tuple[Path, ...]:
        """Modules the filter rejected (read, never parsed)."""
        return self._skipped


class TraceScanner:
    """Collects ``@traces`` annotations from a corpus or a directory."""

    def __init__(self, tracer: Tracer | None = None) -> None:
        """Compose an injected tracer (process default when omitted)."""
        self._tracer = tracer if tracer is not None else Tracer.default()

    def collect_from(self, corpus: Iterable[ParsedModule]) -> list[TraceInfo]:
        """Scan already-parsed modules — no file is read or parsed here."""
        results: list[TraceInfo] = []
        with self._tracer.span("TraceScanner.collect_from"):
            for pm in corpus:
                results.extend(
                    _scan_module_nodes(pm.tree, pm.module_name, str(pm.path))
                )
        return results

    def collect_traces(self, root_dir: str) -> list[TraceInfo]:
        """Walk *root_dir* and scan every module, streaming (no corpus kept)."""
        root = Path(root_dir)
        if not root.is_dir():
            raise ValueError(f"Not a directory: {root_dir}")
        return self.collect_from(_stream_modules(root, AcceptAllFilter()))


def collect_traces(root_dir: str) -> list[TraceInfo]:
    """Walk source tree and collect all @traces annotations via AST.

    Implemented as ``TraceScanner.collect_from`` over an accept-all stream of
    the modules, so it is byte-identical to scanning a ``ModuleCorpus`` loaded
    with the default filter (WI-FLOSS).

    Args:
        root_dir: Root directory to scan (e.g., "src/").

    Returns:
        List of TraceInfo objects, one per requirement ID found.
    """
    return TraceScanner().collect_traces(root_dir)