orchestrata_ultimate.knowledge_graph.class_extractor

Reference › orchestrata_ultimate › knowledge_graph › class_extractor

ClassExtractor turns Python ast class definitions into ClassNode objects (see types_nodes) and is used by ast_analyzer. extract returns every class in a parsed module for a given file path. extract_with_methods returns each class paired with the FunctionNode list of its methods, built by _extract_methods using the class id so methods link to their owner. _extract_class builds one node and _extract_decorators collects decorator names. It operates only on the syntax tree it is handed and touches no graph, Redis or filesystem. The sibling for functions is function_extractor.

At a glance

Class diagram

Dependency diagram

API

Classes

ClassExtractor

Line 17

Extracts class information from Python AST nodes.

Source

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

# mypy: ignore-errors
"""Class extraction from Python AST.

Extracts class definitions, methods, base classes, decorators,
and docstrings from parsed Python source.
"""

import ast
import logging

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.knowledge_graph.types import ClassNode, FunctionNode

logger = logging.getLogger(__name__)


class ClassExtractor:
    """Extracts class information from Python AST nodes."""

    @traces(spec="008", story="US-008-001", reqs=["REQ-008-100"])
    def extract(self, tree: ast.Module, file_path: str) -> list[ClassNode]:
        """Extract all class definitions from an AST."""
        classes = []
        for node in ast.iter_child_nodes(tree):
            if isinstance(node, ast.ClassDef):
                class_node = self._extract_class(node, file_path)
                classes.append(class_node)
        return classes

    @traces(spec="008", story="US-004", reqs=["REQ-008-075"])
    def extract_with_methods(
        self, tree: ast.Module, file_path: str
    ) -> list[tuple[ClassNode, list[FunctionNode]]]:
        """Extract classes with their method details."""
        results = []
        for node in ast.iter_child_nodes(tree):
            if isinstance(node, ast.ClassDef):
                class_node = self._extract_class(node, file_path)
                methods = self._extract_methods(node, file_path, class_node.id)
                results.append((class_node, methods))
        return results

    @traces(spec="008", story="US-004", reqs=["REQ-008-065"])
    def _extract_class(self, node: ast.ClassDef, file_path: str) -> ClassNode:
        """Extract a single class definition."""
        methods = []
        for item in node.body:
            if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
                methods.append(item.name)

        base_classes = []
        for base in node.bases:
            if isinstance(base, ast.Name):
                base_classes.append(base.id)
            elif isinstance(base, ast.Attribute):
                parts = []
                current = base
                while isinstance(current, ast.Attribute):
                    parts.append(current.attr)
                    current = current.value
                if isinstance(current, ast.Name):
                    parts.append(current.id)
                base_classes.append(".".join(reversed(parts)))

        decorators = self._extract_decorators(node.decorator_list)
        docstring = ast.get_docstring(node) or ""

        return ClassNode(
            name=node.name,
            file_id=file_path,
            line_start=node.lineno,
            line_end=node.end_lineno or node.lineno,
            docstring=docstring,
            methods=methods,
            base_classes=base_classes,
            decorators=decorators,
        )

    @traces(
        spec="008",
        story="US-004",
        reqs=["REQ-008-072", "REQ-008-073", "REQ-008-074"],
    )
    def _extract_methods(
        self, class_node: ast.ClassDef, file_path: str, class_id: str
    ) -> list[FunctionNode]:
        """Extract method details from a class."""
        methods = []
        for item in class_node.body:
            if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
                params = [arg.arg for arg in item.args.args]
                return_type = ""
                if item.returns:
                    try:
                        return_type = ast.unparse(item.returns)
                    except Exception as exc:
                        logger.warning(
                            "Failed to unparse method return type for %s: %s",
                            item.name,
                            exc,
                            exc_info=True,
                        )

                decorators = self._extract_decorators(item.decorator_list)
                docstring = ast.get_docstring(item) or ""

                method = FunctionNode(
                    name=item.name,
                    file_id=file_path,
                    class_id=class_id,
                    line_start=item.lineno,
                    line_end=item.end_lineno or item.lineno,
                    signature=f"def {item.name}({', '.join(params)})",
                    parameters=params,
                    return_type=return_type,
                    decorators=decorators,
                    is_async=isinstance(item, ast.AsyncFunctionDef),
                    docstring=docstring,
                )
                methods.append(method)
        return methods

    @traces(spec="008", story="US-004", reqs=["REQ-008-072"])
    def _extract_decorators(self, decorator_list: list) -> list[str]:
        """Extract decorator names from a list of decorator nodes."""
        decorators = []
        for dec in decorator_list:
            if isinstance(dec, ast.Name):
                decorators.append(dec.id)
            elif isinstance(dec, ast.Attribute):
                try:
                    decorators.append(ast.unparse(dec))
                except Exception as exc:
                    logger.warning(
                        "Failed to unparse attribute decorator: %s", exc, exc_info=True
                    )
                    decorators.append(dec.attr)
            elif isinstance(dec, ast.Call):
                if isinstance(dec.func, ast.Name):
                    decorators.append(dec.func.id)
                elif isinstance(dec.func, ast.Attribute):
                    try:
                        decorators.append(ast.unparse(dec.func))
                    except Exception as exc:
                        logger.warning(
                            "Failed to unparse call decorator attribute: %s",
                            exc,
                            exc_info=True,
                        )
                        decorators.append(dec.func.attr)
        return decorators