orchestrata_ultimate.tools.ast_query

Reference › orchestrata_ultimate › tools › ast_query

This module is the AST Query Engine: programmatic access to a serialized AST snapshot (plans/ast-ultimate.json by default) and, optionally, a comparison file against the frozen source. The command line in ast_query_cli is its only importer; tests are tests/test_ast_query.py and tests/tools/test_ast_query_cli.py. The snapshot is produced by ast_pipeline and its format is defined in orchestrata_ast.models.

ASTQuery loads the JSON and builds internal indexes in _build. ClassInfo and MethodInfo are lightweight summaries. Query methods fall into groups:

• Summaries: summary, module_summary, index_stats.
• Lookup: find_class, search_classes, search_methods, find_module, all_classes, classes_by_module.
• Shape: stubs (classes with few methods against their source), largest_classes, most_connected_classes, abstract_classes, pydantic_models, dataclasses, protocols, enums, plus interfaces, type_aliases and exported_classes for JavaScript and TypeScript snapshots.
• Structure: inheritance_tree, module_deps, importers_of.
• Comparison: coverage_by_tier, unmapped_classes, method_diff, which need the comparison file.
• Freshness: staleness and staleness_warning check whether the live source tree is newer than the snapshot.

The helpers _find_project_root and _resolve_path locate the snapshot by walking up from the working directory for plans/ast-ultimate.json or a pyproject.toml, and also accept absolute paths.

At a glance

Class diagram

Dependency diagram

API

Classes

ClassInfo (dataclass)

Decorators: dataclass · Line 66

Lightweight class summary for fast iteration.

MethodInfo (dataclass)

Decorators: dataclass · Line 89

Lightweight method/function summary.

ASTQuery

Line 102

Query engine over serialized Python AST JSON.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/tools/ast_query.py

"""
AST Query Engine — Programmatic access to serialized Python AST.

Loads ``plans/ast-ultimate.json`` (and optionally the orchestrata comparison)
and exposes a rich query API: find classes, list stubs, compute coverage,
trace inheritance, search methods, and more.

Usage::

    from orchestrata_ultimate.tools.ast_query import ASTQuery

    ast = ASTQuery()
    ast.summary()
    ast.stubs()
    ast.find_class("BaseAgent")
    ast.classes_by_module()
"""

from __future__ import annotations

import json
import re
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any

from orchestrata_ultimate.core.paths import resolve_project_root


def _find_project_root() -> Path | None:
    """Walk up from CWD looking for a project with plans/ast-ultimate.json."""
    current = Path.cwd().resolve()
    for parent in [current, *current.parents]:
        if (parent / "plans" / "ast-ultimate.json").exists():
            return parent
        if (parent / "pyproject.toml").exists():
            return parent
    return None


def _resolve_path(path: str | Path) -> Path:
    """Resolve a path relative to the project root, CWD, or as absolute."""
    p = Path(path)
    if p.is_absolute():
        return p

    project_root = _find_project_root()
    candidates = []
    if project_root:
        candidates.append(project_root / p)
    candidates.extend(
        [
            resolve_project_root() / p,
            Path.cwd() / p,
        ]
    )
    for c in candidates:
        if c.exists():
            return c
    return candidates[0]


@dataclass
class ClassInfo:
    """Lightweight class summary for fast iteration."""

    name: str
    qualified_name: str
    module_path: str
    file_path: str
    method_count: int
    base_classes: list[str] = field(default_factory=list)
    is_abstract: bool = False
    is_dataclass: bool = False
    is_pydantic: bool = False
    is_protocol: bool = False
    is_enum: bool = False
    # JS/TS extensions (zero for Python)
    is_interface: bool = False
    is_type_alias: bool = False
    is_exported: bool = False
    type_parameters: str | None = None
    line: int = 0


@dataclass
class MethodInfo:
    """Lightweight method/function summary."""

    name: str
    qualified_name: str
    module_path: str
    is_async: bool = False
    is_abstract: bool = False
    parameters: list[dict] = field(default_factory=list)
    return_type: str | None = None
    call_count: int = 0


class ASTQuery:
    """Query engine over serialized Python AST JSON."""

    def __init__(
        self,
        ast_path: str | Path = "plans/ast-ultimate.json",
        compare_path: str | Path | None = None,
    ):
        self._ast_path = _resolve_path(ast_path)
        self._compare_path = _resolve_path(compare_path) if compare_path else None

        with open(self._ast_path) as f:
            self._data: dict = json.load(f)

        self._compare_data: dict | None = None
        if self._compare_path and self._compare_path.exists():
            with open(self._compare_path) as f:
                self._compare_data = json.load(f)

        # Indexes built lazily
        self._class_index: dict[str, dict] = {}
        self._class_info: dict[str, ClassInfo] = {}
        self._method_index: dict[str, dict] = {}
        self._module_index: dict[str, dict] = {}
        self._built = False

    # ── building indexes ──────────────────────────────────────────────

    def _build(self) -> None:
        if self._built:
            return
        for mod in self._data.get("modules", []):
            mp = mod["path"]
            self._module_index[mp] = mod
            for cls in mod.get("classes", []):
                qn = cls["qualified_name"]
                self._class_index[qn] = cls
                self._class_info[qn] = ClassInfo(
                    name=cls["name"],
                    qualified_name=qn,
                    module_path=cls.get("module_path", mp),
                    file_path=cls.get("file_path", mod.get("file_path", "")),
                    method_count=len(cls.get("methods", [])),
                    base_classes=cls.get("bases") or [],
                    is_abstract=cls.get("is_abstract", False),
                    is_dataclass=cls.get("is_dataclass", False),
                    is_pydantic=cls.get("is_pydantic", False),
                    is_protocol=cls.get("is_protocol", False),
                    is_enum=cls.get("is_enum", False),
                    is_interface=cls.get("is_interface", False),
                    is_type_alias=cls.get("is_type_alias", False),
                    is_exported=cls.get("is_exported", False),
                    type_parameters=cls.get("type_parameters"),
                    line=cls.get("line", 0),
                )
                for meth in cls.get("methods", []):
                    mqn = meth["qualified_name"]
                    self._method_index[mqn] = meth
        self._built = True

    # ── summaries ─────────────────────────────────────────────────────

    def summary(self) -> dict[str, Any]:
        """Return top-level project summary."""
        self._build()
        return {
            "name": self._data.get("name"),
            "root_path": self._data.get("root_path"),
            "modules": len(self._data.get("modules", [])),
            "classes": self._data.get("total_classes", 0),
            "functions": self._data.get("total_functions", 0),
            "methods": self._data.get("total_methods", 0),
            "total_loc": self._data.get("total_loc", 0),
        }

    def module_summary(self) -> list[dict[str, Any]]:
        """Return per-module LOC, class count, function count."""
        self._build()
        result = [
            {
                "path": mod["path"],
                "file_path": mod.get("file_path", ""),
                "line_count": mod.get("line_count", 0),
                "classes": len(mod.get("classes", [])),
                "functions": len(mod.get("functions", [])),
                "imports": len(mod.get("imports", [])),
            }
            for mod in self._data.get("modules", [])
        ]
        return sorted(result, key=lambda m: m["line_count"], reverse=True)

    # ── finding ───────────────────────────────────────────────────────

    def find_class(self, name: str) -> dict | None:
        """Find a class by exact or partial name. Returns full class dict."""
        self._build()
        # Try exact qualified name first
        if name in self._class_index:
            return self._class_index[name]
        # Try short name
        for qn, cls in self._class_index.items():
            if cls["name"] == name:
                return cls
        # Fuzzy: substring match
        matches = [
            cls for qn, cls in self._class_index.items() if name.lower() in qn.lower()
        ]
        if len(matches) == 1:
            return matches[0]
        return None

    def search_classes(
        self, pattern: str, case_sensitive: bool = False
    ) -> list[ClassInfo]:
        """Return ClassInfo for all classes whose name matches a regex."""
        self._build()
        flags = 0 if case_sensitive else re.IGNORECASE
        rx = re.compile(pattern, flags)
        return [
            ci
            for ci in self._class_info.values()
            if rx.search(ci.name) or rx.search(ci.qualified_name)
        ]

    def search_methods(self, pattern: str, case_sensitive: bool = False) -> list[dict]:
        """Return method dicts matching a regex against qualified name."""
        self._build()
        flags = 0 if case_sensitive else re.IGNORECASE
        rx = re.compile(pattern, flags)
        return [
            m for m in self._method_index.values() if rx.search(m["qualified_name"])
        ]

    def find_module(self, path: str) -> dict | None:
        """Return module dict by path (e.g. 'agents.base_agent')."""
        self._build()
        return self._module_index.get(path)

    # ── analysis ──────────────────────────────────────────────────────

    def all_classes(self) -> list[ClassInfo]:
        """Return all ClassInfo summaries."""
        self._build()
        return list(self._class_info.values())

    def stubs(
        self,
        min_methods_in_source: int = 3,
        max_ratio: float = 0.3,
    ) -> list[dict[str, Any]]:
        """Return classes that appear to be stubs (low method count vs source).

        If a comparison AST is loaded, uses the deep-compare readiness data.
        Otherwise heuristically flags classes with few methods.
        """
        self._build()
        stubs = []

        if self._compare_data:
            # Use deep-comparison data. Support both the legacy schema
            # (`ultimate_class`/`orchestrata_class` dicts with a
            # `method_count`, falling back to a `methods` list) and the
            # schema actually produced by
            # `scripts/serialize_codebase_ast.py deep-compare`, where each
            # pair carries `ultimate`/`orchestrata` dicts whose `methods` is
            # already an int count.
            for pair in self._compare_data.get("pairs", []):
                readiness = pair.get("readiness", "")
                if readiness in ("stub", "partial"):
                    uc = pair.get("ultimate_class") or pair.get("ultimate") or {}
                    oc = pair.get("orchestrata_class") or pair.get("orchestrata") or {}
                    u_methods = uc.get("method_count", uc.get("methods"))
                    if not isinstance(u_methods, int):
                        u_methods = len(u_methods or [])
                    o_methods = oc.get("method_count", oc.get("methods"))
                    if not isinstance(o_methods, int):
                        o_methods = len(o_methods or [])
                    method_diff = pair.get("method_diff", {})
                    stubs.append(
                        {
                            "name": uc.get("name", pair.get("ultimate_name", "?")),
                            "module": uc.get("module_path", uc.get("module", "")),
                            "readiness": readiness,
                            "ultimate_methods": u_methods,
                            "orchestrata_methods": o_methods,
                            "ratio": round(u_methods / max(o_methods, 1), 2),
                            "missing_methods": (
                                pair.get("ultimate_only_methods")
                                or pair.get("orchestrata_only_methods")
                                or method_diff.get("only_in_ultimate")
                                or method_diff.get("only_in_orchestrata")
                                or []
                            ),
                        }
                    )
            return sorted(stubs, key=lambda s: s["ratio"])

        # Heuristic: flag classes with very few methods
        for ci in self._class_info.values():
            # Skip enums, protocols, exceptions, abstract base classes
            if ci.is_enum or ci.is_protocol:
                continue
            if ci.method_count < min_methods_in_source and ci.method_count > 0:
                stubs.append(
                    {
                        "name": ci.name,
                        "module": ci.module_path,
                        "readiness": "heuristic",
                        "ultimate_methods": ci.method_count,
                        "orchestrata_methods": "?",
                        "ratio": 0.0,
                        "missing_methods": [],
                    }
                )
        return sorted(stubs, key=lambda s: s["ultimate_methods"])

    def classes_by_module(self) -> dict[str, list[ClassInfo]]:
        """Return classes grouped by top-level module."""
        self._build()
        groups: dict[str, list[ClassInfo]] = defaultdict(list)
        for ci in self._class_info.values():
            top = ci.module_path.split(".")[0] if ci.module_path else "root"
            groups[top].append(ci)
        return dict(groups)

    def largest_classes(self, n: int = 20) -> list[dict[str, Any]]:
        """Return the n classes with the most methods."""
        self._build()
        ranked = sorted(
            self._class_info.values(),
            key=lambda c: c.method_count,
            reverse=True,
        )[:n]
        return [
            {
                "name": ci.name,
                "qualified_name": ci.qualified_name,
                "module": ci.module_path,
                "methods": ci.method_count,
                "line": ci.line,
            }
            for ci in ranked
        ]

    def most_connected_classes(self, n: int = 20) -> list[dict[str, Any]]:
        """Return classes with the most subclasses (via inheritance graph)."""
        self._build()
        ig = self._data.get("inheritance_graph", {})
        # Count how many classes inherit from each class
        child_count: dict[str, int] = defaultdict(int)
        for child, parents in ig.items():
            for parent in parents:
                child_count[parent] += 1
        ranked = sorted(child_count.items(), key=lambda x: x[1], reverse=True)[:n]
        return [{"class": name, "subclasses": count} for name, count in ranked]

    def abstract_classes(self) -> list[ClassInfo]:
        """Return all abstract classes."""
        self._build()
        return [ci for ci in self._class_info.values() if ci.is_abstract]

    def pydantic_models(self) -> list[ClassInfo]:
        """Return all Pydantic BaseModel classes."""
        self._build()
        return [ci for ci in self._class_info.values() if ci.is_pydantic]

    def dataclasses(self) -> list[ClassInfo]:
        """Return all dataclass-annotated classes."""
        self._build()
        return [ci for ci in self._class_info.values() if ci.is_dataclass]

    def protocols(self) -> list[ClassInfo]:
        """Return all Protocol classes."""
        self._build()
        return [ci for ci in self._class_info.values() if ci.is_protocol]

    def enums(self) -> list[ClassInfo]:
        """Return all Enum classes."""
        self._build()
        return [ci for ci in self._class_info.values() if ci.is_enum]

    def interfaces(self) -> list[ClassInfo]:
        """Return all Interface declarations (JS/TS: ``interface Foo { ... }``)."""
        self._build()
        return [ci for ci in self._class_info.values() if ci.is_interface]

    def type_aliases(self) -> list[ClassInfo]:
        """Return all Type alias declarations (JS/TS: ``type Foo = ...``)."""
        self._build()
        return [ci for ci in self._class_info.values() if ci.is_type_alias]

    def exported_classes(self) -> list[ClassInfo]:
        """Return all exported classes/functions/interfaces (JS/TS: ``export class ...``)."""
        self._build()
        return [ci for ci in self._class_info.values() if ci.is_exported]

    def inheritance_tree(self, class_name: str) -> dict[str, Any] | None:
        """Return parents and children for a class."""
        self._build()
        ig = self._data.get("inheritance_graph", {})
        # Find the fully qualified name
        qn = class_name
        if qn not in ig:
            for k in ig:
                if k.endswith(class_name):
                    qn = k
                    break
        if qn not in ig:
            return None

        children = [k for k, parents in ig.items() if qn in parents]
        return {
            "class": qn,
            "parents": ig.get(qn, []),
            "children": children,
            "child_count": len(children),
        }

    def module_deps(self, module_path: str) -> list[str]:
        """Return modules that this module imports."""
        dg: dict[str, list[str]] = self._data.get("dependency_graph", {})
        return dg.get(module_path, [])

    def importers_of(self, module_path: str) -> list[str]:
        """Return modules that import this module."""
        dg = self._data.get("dependency_graph", {})
        return [k for k, deps in dg.items() if module_path in deps]

    # ── comparison helpers (require compare_path) ─────────────────────

    def coverage_by_tier(self) -> dict[str, dict[str, Any]]:
        """Return coverage stats by tier if comparison data is loaded.

        Supports the legacy ``tier_summary`` list schema as well as the
        ``coverage_by_tier`` mapping schema actually produced by the
        current ``scripts/serialize_codebase_ast.py deep-compare`` command.
        """
        if not self._compare_data:
            return {}
        result: dict[str, dict[str, Any]] = {}
        legacy = self._compare_data.get("tier_summary")
        if legacy:
            for tier_entry in legacy:
                tier = str(tier_entry.get("tier", "?"))
                result[tier] = {
                    "packages": tier_entry.get("package_count", 0),
                    "orchestrata_loc": tier_entry.get("orchestrata_loc", 0),
                    "ultimate_loc": tier_entry.get("ultimate_loc", 0),
                    "coverage": tier_entry.get("loc_coverage", 0),
                    "class_coverage": tier_entry.get("class_coverage", 0),
                }
            return result
        for tier, tier_entry in self._compare_data.get("coverage_by_tier", {}).items():
            result[str(tier)] = {
                "packages": len(tier_entry.get("packages", [])),
                "orchestrata_loc": tier_entry.get("orch_loc", 0),
                "ultimate_loc": tier_entry.get("ult_loc", 0),
                "coverage": tier_entry.get("loc_coverage_pct", 0),
                "class_coverage": tier_entry.get("class_coverage_pct", 0),
            }
        return result

    def unmapped_classes(self) -> dict[str, list[str]]:
        """Return orchestrata classes not yet ported and ultimate orphan classes.

        Supports the legacy top-level ``orchestrata_only``/``ultimate_only``
        lists of class dicts as well as the ``unmapped`` mapping produced by
        the current deep-compare command. In that mapping ``ultimate_only``
        is a flat list of class-name strings, but the unmapped orchestrata
        classes are grouped by package under ``orchestrata_by_package``
        (a ``{package: [class_name, ...]}`` mapping) — NOT a flat
        ``orchestrata_only`` list — so they must be flattened out.
        """
        if not self._compare_data:
            return {}
        unmapped = self._compare_data.get("unmapped")
        if unmapped is not None:
            orch = list(unmapped.get("orchestrata_only", []))
            if not orch:
                by_pkg = unmapped.get("orchestrata_by_package", {})
                if isinstance(by_pkg, dict):
                    orch = [name for names in by_pkg.values() for name in names]
            return {
                "orchestrata_only": orch,
                "ultimate_only": list(unmapped.get("ultimate_only", [])),
            }
        return {
            "orchestrata_only": [
                c.get("name", "?")
                for c in self._compare_data.get("orchestrata_only", [])
            ],
            "ultimate_only": [
                c.get("name", "?") for c in self._compare_data.get("ultimate_only", [])
            ],
        }

    def method_diff(self, class_name: str) -> dict[str, Any] | None:
        """Return method-level diff for a class (requires comparison data).

        Supports the legacy schema (``ultimate_class``/``orchestrata_class``
        dicts, flat ``shared_methods``/``param_mismatches`` lists on the
        pair) and the schema produced by the current deep-compare command,
        where each pair nests a ``method_diff`` dict with ``shared``/
        ``only_in_ultimate``/``only_in_orchestrata`` keys.
        """
        if not self._compare_data:
            return None
        for pair in self._compare_data.get("pairs", []):
            uc = pair.get("ultimate_class", {})
            oc = pair.get("orchestrata_class", {})
            u_name = uc.get("name") or pair.get("ultimate_name", "?")
            o_name = oc.get("name") or pair.get("orchestrata_name", "?")
            if u_name != class_name and o_name != class_name:
                continue

            method_diff = pair.get("method_diff")
            if method_diff is None:
                return {
                    "pair_name": f"{u_name} ↔ {o_name}",
                    "readiness": pair.get("readiness", "?"),
                    "ultimate_method_count": uc.get(
                        "method_count", len(uc.get("methods", []))
                    ),
                    "orchestrata_method_count": oc.get(
                        "method_count", len(oc.get("methods", []))
                    ),
                    "shared_methods": pair.get("shared_methods", []),
                    "ultimate_only_methods": pair.get("ultimate_only_methods", []),
                    "orchestrata_only_methods": pair.get(
                        "orchestrata_only_methods", []
                    ),
                    "param_mismatches": pair.get("param_mismatches", []),
                }

            ultimate = pair.get("ultimate", {})
            orchestrata = pair.get("orchestrata", {})
            shared = method_diff.get("shared", [])
            return {
                "pair_name": f"{u_name} ↔ {o_name}",
                "readiness": pair.get("readiness", "?"),
                "ultimate_method_count": ultimate.get("methods", 0),
                "orchestrata_method_count": orchestrata.get("methods", 0),
                "shared_methods": [m.get("name", "?") for m in shared],
                "ultimate_only_methods": method_diff.get("only_in_ultimate", []),
                "orchestrata_only_methods": method_diff.get("only_in_orchestrata", []),
                "param_mismatches": [
                    m.get("name", "?")
                    for m in shared
                    if m.get("param_diffs") or m.get("params_a") != m.get("params_b")
                ],
            }
        return None

    # ── introspection ─────────────────────────────────────────────────

    def index_stats(self) -> dict[str, int]:
        """Return size of internal indexes."""
        self._build()
        return {
            "classes": len(self._class_index),
            "methods": len(self._method_index),
            "modules": len(self._module_index),
        }

    # ── snapshot freshness ────────────────────────────────────────────

    def staleness(self) -> tuple[int, float] | None:
        """Check whether the live source tree is newer than this snapshot.

        Compares the snapshot file's mtime against every ``.py`` file under
        the snapshot's recorded ``root_path``. Returns ``None`` if the
        snapshot appears fresh (no newer source file found) or its source
        root can't be located; otherwise returns
        ``(stale_file_count, snapshot_mtime)``.
        """
        root_path = self._data.get("root_path")
        if not root_path:
            return None
        src_root = _resolve_path(root_path)
        if not src_root.exists():
            return None
        try:
            snapshot_mtime = self._ast_path.stat().st_mtime
        except OSError:
            return None
        stale = 0
        for py_file in src_root.rglob("*.py"):
            try:
                if py_file.stat().st_mtime > snapshot_mtime:
                    stale += 1
            except OSError:
                continue
        if stale == 0:
            return None
        return stale, snapshot_mtime

    def staleness_warning(self) -> str | None:
        """Return a one-line staleness WARNING, or ``None`` if the snapshot is fresh.

        Includes the snapshot's generation date and the number of stale
        (changed-since-snapshot) source files, so a caller can print it to
        stderr without altering its own stdout output format.
        """
        info = self.staleness()
        if info is None:
            return None
        stale, snapshot_mtime = info
        snap_date = datetime.fromtimestamp(snapshot_mtime).strftime("%Y-%m-%d %H:%M")
        return (
            f"WARNING: AST snapshot {self._ast_path} is stale "
            f"(generated {snap_date}); {stale} source file(s) under "
            f"{self._data.get('root_path')} changed since — re-run "
            f"'serialize_codebase_ast.py serialize' for accurate numbers."
        )