orchestrata_tools.test_analyzer.call_extraction

Reference › orchestrata_tools › test_analyzer › call_extraction

This module extracts what one function body does, for the test-analysis pipeline. Everything in it is stateless: given one ast node and the module's import context, it produces a CallProfile. The profile replaced an anonymous tuple of calls, used imports and variable types that was threaded through several BodyMatcher signatures. Its only importer is body_matcher, and tests/test_orchestrata_tools_call_extraction.py covers it.

CallProfile is a frozen dataclass of resolved call names, imported classes used, and inferred var_types; merged_with unions two profiles, with the other profile's variable types winning on conflict. CallExtractor has extract for a node given imported names and classes, plus helpers is_pytest_fixture, fixture_params, decorator_names, get_call_name (dotted name of a call), get_decorator_name, get_annotation_name and collect_orchestrata_imports, which maps aliases of orchestrata_ultimate imports to their original names.

At a glance

Class diagram

Dependency diagram

API

Constants

Classes

CallProfile (dataclass)

Decorators: dataclass(frozen=True) · Line 97

What one function's body does.

CallExtractor

Line 119

Stateless extraction of a :class:CallProfile from one node.

Functions

Source

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

"""Pure per-function call extraction for the test-analysis pipeline.

Everything here is stateless: given one ``ast`` node and the module's
import context, produce a :class:`CallProfile` — the resolved call names,
the imported classes the body uses, and the inferred variable types. The
profile replaces the anonymous ``(calls, used_imports, var_types)`` tuple
that used to thread through five ``BodyMatcher`` signatures.
"""

from __future__ import annotations

import ast
from collections import deque
from collections.abc import Mapping
from dataclasses import dataclass, field
from types import MappingProxyType


def get_call_name(node: ast.expr) -> str | None:
    """Dotted name of a call expression (``a.b.c()`` -> ``"a.b.c"``)."""
    if isinstance(node, ast.Name):
        return node.id
    elif isinstance(node, ast.Attribute):
        parent = get_call_name(node.value)
        if parent:
            return f"{parent}.{node.attr}"
    elif isinstance(node, ast.Subscript):
        return get_call_name(node.value)
    return None


def get_decorator_name(node: ast.expr) -> str:
    """String name of a decorator (``@pytest.fixture()`` -> ``"pytest.fixture"``)."""
    if isinstance(node, ast.Name):
        return node.id
    elif isinstance(node, ast.Attribute):
        parent = get_decorator_name(node.value)
        if parent:
            return f"{parent}.{node.attr}"
    elif isinstance(node, ast.Call):
        return get_decorator_name(node.func)
    return ""


def get_annotation_name(node: ast.expr) -> str | None:
    """String form of a type annotation (subscripts collapse to ``X[...]``)."""
    if isinstance(node, ast.Name):
        return node.id
    elif isinstance(node, ast.Attribute):
        parent = get_annotation_name(node.value)
        if parent:
            return f"{parent}.{node.attr}"
    elif isinstance(node, ast.Subscript):
        parent = get_annotation_name(node.value)
        if parent:
            return f"{parent}[...]"
    return None


def collect_orchestrata_imports(tree: ast.AST) -> tuple[dict[str, str], set[str]]:
    """(alias -> original name, original names) for orchestrata_ultimate imports.

    Import statements can only appear in statement position, so the walk
    descends into statement nodes only (skipping expression subtrees, the
    bulk of the tree). Emission order equals ``ast.walk``'s BFS restricted
    to statements, preserving the original dict-overwrite order.
    """
    imported_names: dict[str, str] = {}
    imported_classes: set[str] = set()
    queue: deque[ast.AST] = deque([tree])
    while queue:
        node = queue.popleft()
        for child in ast.iter_child_nodes(node):
            if isinstance(child, (ast.stmt, ast.ExceptHandler, ast.match_case)):
                queue.append(child)
        if isinstance(node, ast.ImportFrom):
            module_name = node.module or ""
            if module_name.startswith("orchestrata_ultimate"):
                for alias in node.names:
                    name = alias.asname or alias.name
                    imported_names[name] = alias.name
                    imported_classes.add(alias.name)
        elif isinstance(node, ast.Import):
            for alias in node.names:
                parts = alias.name.split(".")
                if parts[0] == "orchestrata_ultimate":
                    name = alias.asname or parts[-1]
                    imported_names[name] = parts[-1]
                    imported_classes.add(parts[-1])
    return imported_names, imported_classes


_EMPTY_TYPES: Mapping[str, str] = MappingProxyType({})


@dataclass(frozen=True)
class CallProfile:
    """What one function's body does.

    Immutable on purpose: fixture profiles are shared across every test
    that uses the fixture, so in-place mutation was an aliasing bug
    waiting to happen. ``var_types`` maps variable name -> class name.
    """

    calls: frozenset[str] = frozenset()
    used_imports: frozenset[str] = frozenset()
    var_types: Mapping[str, str] = field(default=_EMPTY_TYPES)

    def merged_with(self, other: CallProfile) -> CallProfile:
        """Union of both profiles; *other*'s var_types win on conflict
        (same precedence as the original ``dict.update`` merge order)."""
        return CallProfile(
            calls=self.calls | other.calls,
            used_imports=self.used_imports | other.used_imports,
            var_types=MappingProxyType({**self.var_types, **other.var_types}),
        )


class CallExtractor:
    """Stateless extraction of a :class:`CallProfile` from one node."""

    _PRIMITIVES = {"str", "int", "float", "bool", "list", "dict", "set", "tuple"}

    def is_pytest_fixture(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> bool:
        """True if the function carries a ``@pytest.fixture`` decorator."""
        for decorator in node.decorator_list:
            name = get_decorator_name(decorator)
            if name == "pytest.fixture" or name == "fixture":
                return True
        return False

    def fixture_params(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> list[str]:
        """Parameter names that are likely pytest fixtures.

        Skips ``self``/``cls`` and primitive-annotated parameters.
        """
        params = []
        for arg in node.args.args + node.args.posonlyargs + node.args.kwonlyargs:
            name = arg.arg
            if name in ("self", "cls"):
                continue
            if arg.annotation:
                ann_name = get_annotation_name(arg.annotation)
                if ann_name and ann_name.split("[")[0].strip() in self._PRIMITIVES:
                    continue
            params.append(name)
        return params

    def decorator_names(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> set[str]:
        """Resolved names of the node's own decorators (call or bare)."""
        names: set[str] = set()
        for d in node.decorator_list:
            n = get_call_name(d.func if isinstance(d, ast.Call) else d)
            if n:
                names.add(n)
        return names

    def extract(
        self,
        node: ast.AST,
        imported_names: dict[str, str],
        imported_classes: set[str],
        local_classes: set[str] | None = None,
    ) -> CallProfile:
        """Extract calls + used imported classes + inferred var types.

        ``obj = ImportedClass(); obj.method()`` resolves to
        ``ImportedClass.method``.
        """
        local_classes = local_classes or set()
        all_classes = imported_classes | local_classes
        calls: set[str] = set()
        used_imports: set[str] = set()
        var_types: dict[str, str] = {}

        for child in ast.walk(node):
            # Track assignments: servicer = CXPServicer()
            if isinstance(child, ast.Assign):
                for target in child.targets:
                    if isinstance(target, ast.Name) and isinstance(
                        child.value, ast.Call
                    ):
                        func_name = get_call_name(child.value.func)
                        if func_name and func_name in imported_names:
                            var_types[target.id] = imported_names[func_name]
                            used_imports.add(imported_names[func_name])
                        elif func_name and func_name in all_classes:
                            var_types[target.id] = func_name
                            used_imports.add(func_name)
            # Track calls
            if isinstance(child, ast.Call):
                name = get_call_name(child.func)
                if name:
                    parts = name.split(".")
                    # Resolve variable to class:
                    # servicer.send_packet -> CXPServicer.send_packet
                    if len(parts) >= 2 and parts[0] in var_types:
                        resolved = ".".join([var_types[parts[0]]] + parts[1:])
                        calls.add(resolved)
                        used_imports.add(var_types[parts[0]])
                    else:
                        calls.add(name)
                        if parts[0] in imported_names:
                            used_imports.add(imported_names[parts[0]])
                        elif parts[0] in all_classes:
                            used_imports.add(parts[0])

        return CallProfile(
            calls=frozenset(calls),
            used_imports=frozenset(used_imports),
            var_types=MappingProxyType(var_types),
        )