orchestrata_ultimate.knowledge_graph.function_extractor

Reference › orchestrata_ultimate › knowledge_graph › function_extractor

FunctionExtractor turns Python ast function definitions into FunctionNode objects (see types_nodes); it is traced to Spec 008 REQ-008-017 and is used by ast_analyzer. extract returns the top-level functions of a parsed module and extract_all also includes methods. Per function, _extract_function builds the node from helpers that collect parameter names (_extract_params), the return annotation (_extract_return_type), the names of called functions in the body (_extract_calls) and decorator names (_extract_decorators). It works purely on the in-memory syntax tree and the file path it is given; it touches no graph, Redis or filesystem.

At a glance

Class diagram

Dependency diagram

API

Classes

FunctionExtractor

Decorators: traces(spec='008', story='US-004', reqs=['REQ-008-017']) · Line 17

Extracts function information from Python AST nodes.

Source

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

"""Function extraction from Python AST.

Extracts function definitions, parameters, return types,
calls, and decorators from parsed Python source.
"""

import ast
import logging

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

logger = logging.getLogger(__name__)


@traces(spec="008", story="US-004", reqs=["REQ-008-017"])
class FunctionExtractor:
    """Extracts function information from Python AST nodes."""

    def extract(self, tree: ast.Module, file_path: str) -> list[FunctionNode]:
        """Extract all top-level function definitions from an AST."""
        functions = []

        # Identify functions that are methods (inside classes)
        class_bodies: set[int] = {
            id(item)
            for node in ast.iter_child_nodes(tree)
            if isinstance(node, ast.ClassDef)
            for item in node.body
        }

        for node in ast.iter_child_nodes(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                if id(node) not in class_bodies:
                    func_node = self._extract_function(node, file_path)
                    functions.append(func_node)
        return functions

    def extract_all(self, tree: ast.Module, file_path: str) -> list[FunctionNode]:
        """Extract ALL functions including methods."""
        functions = []
        for node in ast.walk(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                func_node = self._extract_function(node, file_path)
                functions.append(func_node)
        return functions

    def _extract_function(
        self, node: ast.FunctionDef | ast.AsyncFunctionDef, file_path: str
    ) -> FunctionNode:
        """Extract a single function definition."""
        params = self._extract_params(node)
        return_type = self._extract_return_type(node)
        decorators = self._extract_decorators(node.decorator_list)
        calls = self._extract_calls(node)
        docstring = ast.get_docstring(node) or ""

        return FunctionNode(
            name=node.name,
            file_id=file_path,
            line_start=node.lineno,
            line_end=node.end_lineno or node.lineno,
            signature=f"def {node.name}({', '.join(params)})",
            parameters=params,
            return_type=return_type,
            decorators=decorators,
            calls=list(set(calls)),
            is_async=isinstance(node, ast.AsyncFunctionDef),
            docstring=docstring,
        )

    def _extract_params(
        self, node: ast.FunctionDef | ast.AsyncFunctionDef
    ) -> list[str]:
        """Extract parameter names."""
        params = [arg.arg for arg in node.args.args]
        if node.args.vararg:
            params.append(f"*{node.args.vararg.arg}")
        if node.args.kwarg:
            params.append(f"**{node.args.kwarg.arg}")
        return params

    def _extract_return_type(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> str:
        """Extract return type annotation."""
        if node.returns:
            try:
                return ast.unparse(node.returns)
            except Exception as exc:
                logger.warning(
                    "Failed to unparse return type annotation for %s: %s",
                    node.name,
                    exc,
                    exc_info=True,
                )
        return ""

    def _extract_calls(self, node: ast.AST) -> list[str]:
        """Extract function call names from a function body."""
        calls = []
        for child in ast.walk(node):
            if isinstance(child, ast.Call):
                if isinstance(child.func, ast.Name):
                    calls.append(child.func.id)
                elif isinstance(child.func, ast.Attribute):
                    calls.append(child.func.attr)
        return calls

    def _extract_decorators(self, decorator_list: list) -> list[str]:
        """Extract decorator names."""
        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)
        return decorators