orchestrata_tools.parser.indexer

Reference › orchestrata_tools › parser › indexer

ASTIndex is the lookup layer that most of orchestrata_tools reads. It is built once over an ASTCodebase (see ast_models) and fills dictionaries in _build: classes by name and by module, methods by qualified name, name and class, functions by name and module, modules by path, plus subclass and caller graphs. Importers include the complexity analyzer, the diagram generators, the project context builder, the test matcher and the reviewers. Tests are tests/test_orchestrata_tools_matcher.py, tests/test_orchestrata_tools_diagrams_deps.py, tests/test_diagram_to_code_roundtrip.py and tests/test_orchestrata_tools_reviewer_query_indexer.py.

Its methods fall into groups:

• Lookup: find_class, find_classes_in_module, find_method, find_methods_by_name, find_methods_in_class, find_functions_by_name, find_functions_in_module and find_module.
• Relations: find_subclasses, find_superclass_chain and find_callers.
• Selections: public_methods (the methods that should have direct tests), abstract_classes, protocol_classes, dataclass_classes, pydantic_classes, god_classes (more than 20 methods by default), async_methods, property_methods, methods_without_docstrings and classes_without_docstrings.
• Rankings: largest_modules and largest_classes.

At a glance

Class diagram

Dependency diagram

API

Classes

ASTIndex

Line 18

Fast lookup indexes for an ASTCodebase.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_tools/parser/indexer.py

"""Build fast lookup indexes from ASTCodebase."""

from __future__ import annotations

from collections import defaultdict
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from orchestrata_tools.models.ast_models import (
        ASTCodebase,
        ClassInfo,
        FunctionInfo,
        MethodInfo,
        ModuleInfo,
    )


class ASTIndex:
    """Fast lookup indexes for an ASTCodebase.

    Provides O(1) or O(n) lookups for classes, methods, functions,
    inheritance relationships, and call graphs.
    """

    def __init__(self, codebase: ASTCodebase) -> None:
        self.codebase = codebase

        # Primary indexes
        self._classes_by_name: dict[str, ClassInfo] = {}
        self._classes_by_module: dict[str, list[ClassInfo]] = defaultdict(list)
        self._methods_by_qualified_name: dict[str, MethodInfo] = {}
        self._methods_by_name: dict[str, list[MethodInfo]] = defaultdict(list)
        self._methods_by_class: dict[str, list[MethodInfo]] = defaultdict(list)
        self._functions_by_name: dict[str, list[FunctionInfo]] = defaultdict(list)
        self._functions_by_module: dict[str, list[FunctionInfo]] = defaultdict(list)
        self._modules_by_path: dict[str, ModuleInfo] = {}

        # Graph indexes
        self._subclasses: dict[str, list[ClassInfo]] = defaultdict(list)
        self._callers: dict[str, list[MethodInfo]] = defaultdict(list)

        self._build()

    def _build(self) -> None:
        for module in self.codebase.modules:
            self._modules_by_path[module.path] = module
            # Make copies to avoid mutating original module/class objects
            self._classes_by_module[module.path] = list(module.classes)
            self._functions_by_module[module.path] = list(module.functions)

            for func in module.functions:
                self._functions_by_name[func.name].append(func)

            for cls in module.classes:
                self._classes_by_name[cls.name] = cls
                self._methods_by_class[cls.name] = list(cls.methods)

                for method in cls.methods:
                    self._methods_by_qualified_name[method.qualified_name] = method
                    self._methods_by_name[method.name].append(method)
                    self._methods_by_class[cls.name].append(method)

                    # Build caller graph
                    for called in method.calls:
                        self._callers[called].append(method)

                # Build subclass index
                for base in cls.bases:
                    self._subclasses[base].append(cls)

    # ------------------------------------------------------------------
    # Class lookups
    # ------------------------------------------------------------------

    def find_class(self, name: str) -> ClassInfo | None:
        return self._classes_by_name.get(name)

    def find_classes_in_module(self, module_path: str) -> list[ClassInfo]:
        return self._classes_by_module.get(module_path, [])

    def find_subclasses(self, base_class: str) -> list[ClassInfo]:
        return self._subclasses.get(base_class, [])

    def find_superclass_chain(self, class_name: str) -> list[str]:
        chain = []
        cls = self.find_class(class_name)
        while cls and cls.bases:
            base_name = cls.bases[0]
            chain.append(base_name)
            cls = self.find_class(base_name)
        return chain

    # ------------------------------------------------------------------
    # Method lookups
    # ------------------------------------------------------------------

    def find_method(self, qualified_name: str) -> MethodInfo | None:
        return self._methods_by_qualified_name.get(qualified_name)

    def find_methods_by_name(self, name: str) -> list[MethodInfo]:
        return self._methods_by_name.get(name, [])

    def find_methods_in_class(self, class_name: str) -> list[MethodInfo]:
        return self._methods_by_class.get(class_name, [])

    def find_callers(self, method_name: str) -> list[MethodInfo]:
        return self._callers.get(method_name, [])

    # ------------------------------------------------------------------
    # Function lookups
    # ------------------------------------------------------------------

    def find_functions_by_name(self, name: str) -> list[FunctionInfo]:
        return self._functions_by_name.get(name, [])

    def find_functions_in_module(self, module_path: str) -> list[FunctionInfo]:
        return self._functions_by_module.get(module_path, [])

    # ------------------------------------------------------------------
    # Module lookups
    # ------------------------------------------------------------------

    def find_module(self, path: str) -> ModuleInfo | None:
        return self._modules_by_path.get(path)

    # ------------------------------------------------------------------
    # Filters
    # ------------------------------------------------------------------

    def public_methods(self) -> list[MethodInfo]:
        """Return public methods that should have direct tests.

        Excludes: private/dunder methods, properties, and protocol methods.
        """
        result = []
        for m in self.codebase.all_methods:
            if not m.is_public:
                continue
            if m.is_property:
                continue
            cls = self._classes_by_name.get(m.class_name)
            if cls and cls.is_protocol:
                continue
            result.append(m)
        return result

    def abstract_classes(self) -> list[ClassInfo]:
        return [c for c in self.codebase.all_classes if c.is_abstract]

    def protocol_classes(self) -> list[ClassInfo]:
        return [c for c in self.codebase.all_classes if c.is_protocol]

    def dataclass_classes(self) -> list[ClassInfo]:
        return [c for c in self.codebase.all_classes if c.is_dataclass]

    def pydantic_classes(self) -> list[ClassInfo]:
        return [c for c in self.codebase.all_classes if c.is_pydantic]

    def god_classes(self, method_threshold: int = 20) -> list[ClassInfo]:
        return [
            c for c in self.codebase.all_classes if c.method_count > method_threshold
        ]

    def async_methods(self) -> list[MethodInfo]:
        return [m for m in self.codebase.all_methods if m.is_async]

    def property_methods(self) -> list[MethodInfo]:
        return [m for m in self.codebase.all_methods if m.is_property]

    def methods_without_docstrings(self) -> list[MethodInfo]:
        return [m for m in self.codebase.all_methods if not m.docstring]

    def classes_without_docstrings(self) -> list[ClassInfo]:
        return [c for c in self.codebase.all_classes if not c.docstring]

    def largest_modules(self, n: int = 10) -> list[tuple[str, int, int]]:
        """Return top N modules by (path, line_count, class_count)."""
        sorted_modules = sorted(
            self.codebase.modules,
            key=lambda m: m.line_count,
            reverse=True,
        )
        return [(m.path, m.line_count, m.class_count) for m in sorted_modules[:n]]

    def largest_classes(self, n: int = 10) -> list[tuple[str, int]]:
        """Return top N classes by method count."""
        sorted_classes = sorted(
            self.codebase.all_classes,
            key=lambda c: c.method_count,
            reverse=True,
        )
        return [(c.qualified_name, c.method_count) for c in sorted_classes[:n]]