orchestrata_ultimate.refactor.codemod._cig_builder_ast

Reference › orchestrata_ultimate › refactor › codemod › _cig_builder_ast

The only backend that builds the CodeIntelligenceGraph, by walking CPython ast trees. workspace is its importer.

build_cig_ast(sources, parallel, trees) takes repo-relative paths mapped to source text, optionally with already-parsed trees so no file is parsed twice. For each module an _AstExtractor walks the tree and records module, class, function and import nodes plus call, inheritance and containment edges into a sink. In parallel mode, _ModuleRecorder records an op log per module on a thread pool and _CIGBuilder.replay applies the logs serially, so the output matches the serial build; logs are freed as they are replayed to limit peak memory. Imports are resolved in a second pass (_resolve_imports) once every module is known: _ImportResolver maps an import, absolute or relative, to the registered module it names, and one IMPORTS edge joins each module pair regardless of how many statements import it; the resolved path is stored on the Import node. build produces the graph with its CSR index. The sink protocol _CIGSink is shared by the builder and recorder. The design history is in the refactor package narrative.

At a glance

Source src/orchestrata_ultimate/refactor/codemod/_cig_builder_ast.py
Lines of code 838
Classes 5
Functions 7
Imports (internal) orchestrata_ultimate.refactor.codemod._cig, orchestrata_ultimate.refactor.codemod._csr_index, orchestrata_ultimate.refactor.codemod._types
Imported by orchestrata_ultimate.refactor.codemod.workspace
Tests that import it tests/refactor/codemod/test_cig_builder_ast.py, tests/refactor/codemod/test_review_checks.py, tests/refactor/codemod/test_workspace_load.py

Class diagram

Class diagram orchestrata_ultimate.refactor.codemod._cig_builder_ast

Dependency diagram

Dependency diagram orchestrata_ultimate.refactor.codemod._cig_builder_ast

API

Classes

_CIGSink (protocol)

Bases: Protocol · Line 42

Shared node/edge-recording interface of _CIGBuilder and _ModuleRecorder.

Method Signature Summary
sink_node sink_node(key: str, node_type: str, name: str, qualified_name: str, path: str, line: int, properties: dict[str, Any] | None = None) -> int
sink_edge sink_edge(source_id: int, target_id: int, edge_type: CIGEdgeType) -> None
sink_edge_to_key sink_edge_to_key(source_id: int, target_key: str, edge_type: CIGEdgeType) -> None
sink_edge_from_key sink_edge_from_key(target_key: str, dest_id: int, edge_type: CIGEdgeType) -> None
sink_call_edge sink_call_edge(source_id: int, callee_name: str) -> None
sink_import sink_import(module_id: int, import_id: int, module_name: str, imported_name: str | None, level: int) -> None

_CIGBuilder

Line 170

Accumulates nodes and edges from parsed modules, then builds the CIG.

Method Signature Summary
__init__ __init__() -> None
_alloc_id _alloc_id() -> int
_add_node _add_node(key: str, node_type: str, name: str, qualified_name: str, path: str, line: int, properties: dict[str, Any] | None = None) -> int
_add_edge _add_edge(source_id: int, target_id: int, edge_type: CIGEdgeType) -> None
sink_node sink_node(key: str, node_type: str, name: str, qualified_name: str, path: str, line: int, properties: dict[str, Any] | None = None) -> int
sink_edge sink_edge(source_id: int, target_id: int, edge_type: CIGEdgeType) -> None
sink_edge_to_key sink_edge_to_key(source_id: int, target_key: str, edge_type: CIGEdgeType) -> None Add an edge to the node for target_key, if already known.
sink_edge_from_key sink_edge_from_key(target_key: str, dest_id: int, edge_type: CIGEdgeType) -> None Add an edge from the node for target_key, if already known.
sink_call_edge sink_call_edge(source_id: int, callee_name: str) -> None Add a CALLS edge to the first-inserted node matching callee_name.
sink_import sink_import(module_id: int, import_id: int, module_name: str, imported_name: str | None, level: int) -> None Record a raw import; its IMPORTS edge is resolved in build.
_resolve_imports _resolve_imports() -> None Second pass: link each import to its module once all are known.
replay replay(ops: list[tuple]) -> None Apply one module's op log, resolving edges against current state.
build build() -> CodeIntelligenceGraph Construct the CIG with CSR index from accumulated nodes and edges.

_ModuleRecorder

Line 361

Sink that records a module's nodes/edges as an op log (thread-local).

Method Signature Summary
__init__ __init__() -> None
sink_node sink_node(key: str, node_type: str, name: str, qualified_name: str, path: str, line: int, properties: dict[str, Any] | None = None) -> int
sink_edge sink_edge(source_id: int, target_id: int, edge_type: CIGEdgeType) -> None
sink_edge_to_key sink_edge_to_key(source_id: int, target_key: str, edge_type: CIGEdgeType) -> None
sink_edge_from_key sink_edge_from_key(target_key: str, dest_id: int, edge_type: CIGEdgeType) -> None
sink_call_edge sink_call_edge(source_id: int, callee_name: str) -> None
sink_import sink_import(module_id: int, import_id: int, module_name: str, imported_name: str | None, level: int) -> None

_AstExtractor

Line 416

Mirror of _cig_builder._ModuleVisitor over CPython ast trees.

Method Signature Summary
__init__ __init__(module_id: int, path: str, sink: _CIGSink) -> None
walk walk(root: ast.AST) -> None Iterative pre-order walk in CST-matching child order.
_class_children staticmethod _class_children(node: ast.ClassDef) -> list[ast.AST]
_func_children staticmethod _func_children(node: ast.FunctionDef | ast.AsyncFunctionDef) -> list[ast.AST]
_visit_class _visit_class(node: ast.ClassDef) -> None
_visit_func _visit_func(node: ast.FunctionDef | ast.AsyncFunctionDef) -> None
_visit_import _visit_import(node: ast.Import) -> None
_visit_import_from _visit_import_from(node: ast.ImportFrom) -> None
_record_import _record_import(module_name: str, imported_name: str | None, level: int, line: int) -> None
_visit_call _visit_call(node: ast.Call) -> None

_ImportResolver

Line 733

Map an import to the registered module it names.

Method Signature Summary
__init__ __init__(paths: set[str]) -> None
_package_root _package_root(init_path: str) -> str Directory containing the outermost package above init_path.
_lookup _lookup(root: str, dotted: list[str]) -> str | None
resolve resolve(importer: str, module_name: str, imported_name: str | None, level: int) -> str | None Path of the module an import names, or None (external or unknown).
_package _package(root: str) -> str | None
_relative_roots _relative_roots(importer: str, level: int) -> list[str] Package directory a level-dot relative import starts from.
_roots_for _roots_for(importer: str) -> list[str]

Functions

Function Signature Summary
build_cig_ast build_cig_ast(sources: Mapping[str, str], parallel: bool = False, trees: dict[str, ast.Module] | None = None) -> CodeIntelligenceGraph Build the CIG from source text via CPython's ast.
_record_module_ast _record_module_ast(item: tuple[str, Mapping[str, str], ast.Module | None]) -> list[tuple] Parse and extract one module on a worker thread, returning its ops.
_extract_module_ast _extract_module_ast(path: str, sources: Mapping[str, str], sink: _CIGSink, tree: ast.Module | None = None) -> None Extract nodes and edges from one module into a sink.
_extract_name _extract_name(node: ast.AST | None) -> str | None Leftmost simple name (mirrors _cig_builder._extract_name).
_extract_dotted_name _extract_dotted_name(node: ast.AST | None) -> str | None Dotted name, e.g. 'foo.bar.baz' (None for anything else).
_extract_annotation _extract_annotation(node: ast.AST | None) -> str | None Type annotation as string.
_extract_call_name _extract_call_name(node: ast.AST) -> str | None Name of the called function (rightmost for attributes).

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/refactor/codemod/_cig_builder_ast.py

"""AST-based Code Intelligence Graph extraction — the sole CIG backend.

Builds the CIG from CPython's ``ast`` (C-implemented parser, ~10x faster
than a LibCST visitor walk would be) over the module's source text. The
forest holds ``LosslessModule`` objects for round-trip fidelity
(REQ-110-002); the CIG is an index over names and structure, which ``ast``
provides with exact line numbers.

Also owns the backend-agnostic op-log sink/builder plumbing
(``_CIGBuilder``, ``_ModuleRecorder``) formerly shared with the deleted
LibCST-walking ``_cig_builder.py``: per-module extraction (here, an
``_AstExtractor`` walk) is the expensive part and has no inter-module
dependencies, so with ``parallel=True`` each module is visited on a thread
that records an op log into a private ``_ModuleRecorder``. A serial replay
phase then applies the logs in forest order, so cross-module edge
resolution (IMPORTS, CALLS, TESTED_BY) observes exactly the keys the
serial path would have observed at that point — output is identical to
the serial build.

Complex type annotations (anything beyond a plain dotted name) are stored
as ``ast.unparse`` output (e.g. ``Dict[str, int]``); this affects only the
``return_type`` / ``type_annotation`` properties.
"""

from __future__ import annotations

import ast
import logging
import os
import sys
from collections.abc import Mapping
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Protocol

from orchestrata_ultimate.refactor.codemod._cig import CodeIntelligenceGraph
from orchestrata_ultimate.refactor.codemod._csr_index import _CSRIndex
from orchestrata_ultimate.refactor.codemod._types import CIGEdgeType, CIGNode

logger = logging.getLogger(__name__)


class _CIGSink(Protocol):
    """Shared node/edge-recording interface of ``_CIGBuilder`` and ``_ModuleRecorder``.

    ``_AstExtractor`` is used against either implementation: ``_CIGBuilder``
    directly in the serial path, ``_ModuleRecorder`` (an op-log recorder) in
    the parallel path.
    """

    def sink_node(
        self,
        key: str,
        node_type: str,
        name: str,
        qualified_name: str,
        path: str,
        line: int,
        properties: dict[str, Any] | None = None,
    ) -> int: ...

    def sink_edge(
        self, source_id: int, target_id: int, edge_type: CIGEdgeType
    ) -> None: ...

    def sink_edge_to_key(
        self, source_id: int, target_key: str, edge_type: CIGEdgeType
    ) -> None: ...

    def sink_edge_from_key(
        self, target_key: str, dest_id: int, edge_type: CIGEdgeType
    ) -> None: ...

    def sink_call_edge(self, source_id: int, callee_name: str) -> None: ...

    def sink_import(
        self,
        module_id: int,
        import_id: int,
        module_name: str,
        imported_name: str | None,
        level: int,
    ) -> None: ...


def build_cig_ast(
    sources: Mapping[str, str],
    parallel: bool = False,
    trees: dict[str, ast.Module] | None = None,
) -> CodeIntelligenceGraph:
    """Build the CIG from source text via CPython's ast.

    Args:
        sources: Mapping of repo-relative paths to source text. Text is read
            only for paths without a pre-parsed tree, so a lazily decoding
            mapping never holds a second copy of every file.
        parallel: If True, extract per-module op logs on a thread pool,
            then replay serially. Output is identical to the serial build.
        trees: Optional mapping of repo-relative paths to ALREADY-PARSED
            ``ast.Module`` objects. Callers that hold a parsed tree (the
            codemod workspace's forest parses every module once for its own
            lossless representation) pass it here so the file is not parsed
            a second time — measured 2026-08-11 at ~8.7s of a full-tree
            serial ``CodemodWorkspace.load``. Paths missing from ``trees``
            fall back to parsing ``sources[path]``.

    Returns:
        Fully constructed CodeIntelligenceGraph with CSR index.
    """
    builder = _CIGBuilder()
    if parallel and len(sources) > 1:
        # The text is only read for a path with no pre-parsed tree, so a
        # ``Mapping`` that decodes on demand never materialises the text.
        items = [(path, sources, (trees or {}).get(path)) for path in sources]
        max_workers = os.cpu_count() or 4
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            logs = list(executor.map(_record_module_ast, items))
        # Free each op log as it is replayed: the logs and the finished graph
        # would otherwise be live at once (peak RSS, REQ-110-004).
        for index, ops in enumerate(logs):
            logs[index] = []
            builder.replay(ops)
    else:
        for path in sources:
            _extract_module_ast(path, sources, builder, tree=(trees or {}).get(path))
    return builder.build()


def _record_module_ast(
    item: tuple[str, Mapping[str, str], ast.Module | None],
) -> list[tuple]:
    """Parse and extract one module on a worker thread, returning its ops."""
    path, sources, tree = item
    recorder = _ModuleRecorder()
    _extract_module_ast(path, sources, recorder, tree=tree)
    return recorder.ops


def _extract_module_ast(
    path: str,
    sources: Mapping[str, str],
    sink: _CIGSink,
    tree: ast.Module | None = None,
) -> None:
    """Extract nodes and edges from one module into a sink.

    ``sources[path]`` is read only when ``tree`` is None (parse fallback).
    """
    if tree is None:
        try:
            tree = ast.parse(sources[path])
        except SyntaxError as e:
            logger.warning(
                "ast.parse failed for %s: %s; skipping CIG extraction", path, e
            )
            return

    module_id = sink.sink_node(
        key=f"module:{path}",
        node_type="Module",
        name=path,
        qualified_name=path,
        path=path,
        line=1,
        properties=None,
    )
    extractor = _AstExtractor(module_id=module_id, path=path, sink=sink)
    extractor.walk(tree)


class _CIGBuilder:
    """Accumulates nodes and edges from parsed modules, then builds the CIG.

    Also serves as the live sink for ``_AstExtractor`` in the serial path.
    """

    def __init__(self) -> None:
        self._nodes: dict[int, CIGNode] = {}
        self._key_to_id: dict[str, int] = {}
        # Maps a key's final segment (the simple name after the last ':') to
        # keys in first-insertion order, so CALLS resolution is O(1) instead
        # of scanning _key_to_id for every Call node. Keys are stored (not
        # ids) because a duplicate key overwrites _key_to_id in place, and
        # the lookup must mirror that.
        self._name_index: dict[str, list[str]] = {}
        self._next_id: int = 0
        self._edges_list: list[tuple[int, int]] = []
        self._edges_by_type: dict[CIGEdgeType, list[tuple[int, int]]] = {
            et: [] for et in CIGEdgeType
        }
        # Raw import records, resolved to IMPORTS edges in build() once every
        # module is registered (an importer may sort before its importee).
        self._pending_imports: list[tuple[int, int, str, str | None, int]] = []

    def _alloc_id(self) -> int:
        nid = self._next_id
        self._next_id += 1
        return nid

    def _add_node(
        self,
        key: str,
        node_type: str,
        name: str,
        qualified_name: str,
        path: str,
        line: int,
        properties: dict[str, Any] | None = None,
    ) -> int:
        nid = self._alloc_id()
        self._nodes[nid] = CIGNode(
            node_id=nid,
            node_type=node_type,
            name=name,
            qualified_name=qualified_name,
            path=path,
            line=line,
            properties=properties or {},
        )
        if key not in self._key_to_id:
            self._name_index.setdefault(key.rsplit(":", 1)[-1], []).append(key)
        self._key_to_id[key] = nid
        return nid

    def _add_edge(self, source_id: int, target_id: int, edge_type: CIGEdgeType) -> None:
        # One tuple object serves both the insertion-ordered CSR input and the
        # per-type list (a ``CIGEdge`` object per edge doubled this footprint).
        edge = (source_id, target_id)
        self._edges_list.append(edge)
        self._edges_by_type[edge_type].append(edge)

    # --- sink interface (serial path) ---

    def sink_node(
        self,
        key: str,
        node_type: str,
        name: str,
        qualified_name: str,
        path: str,
        line: int,
        properties: dict[str, Any] | None = None,
    ) -> int:
        return self._add_node(
            key, node_type, name, qualified_name, path, line, properties
        )

    def sink_edge(self, source_id: int, target_id: int, edge_type: CIGEdgeType) -> None:
        self._add_edge(source_id, target_id, edge_type)

    def sink_edge_to_key(
        self, source_id: int, target_key: str, edge_type: CIGEdgeType
    ) -> None:
        """Add an edge to the node for target_key, if already known."""
        target_id = self._key_to_id.get(target_key)
        if target_id is not None:
            self._add_edge(source_id, target_id, edge_type)

    def sink_edge_from_key(
        self, target_key: str, dest_id: int, edge_type: CIGEdgeType
    ) -> None:
        """Add an edge from the node for target_key, if already known."""
        source_id = self._key_to_id.get(target_key)
        if source_id is not None:
            self._add_edge(source_id, dest_id, edge_type)

    def sink_call_edge(self, source_id: int, callee_name: str) -> None:
        """Add a CALLS edge to the first-inserted node matching callee_name."""
        candidates = self._name_index.get(callee_name)
        if not candidates:
            candidates = self._name_index.get(callee_name.split(".")[-1])
        if candidates:
            self._add_edge(source_id, self._key_to_id[candidates[0]], CIGEdgeType.CALLS)

    def sink_import(
        self,
        module_id: int,
        import_id: int,
        module_name: str,
        imported_name: str | None,
        level: int,
    ) -> None:
        """Record a raw import; its IMPORTS edge is resolved in ``build``."""
        self._pending_imports.append(
            (module_id, import_id, module_name, imported_name, level)
        )

    def _resolve_imports(self) -> None:
        """Second pass: link each import to its module once all are known.

        One IMPORTS edge per (importing module, imported module) pair. The
        resolved path is stored on the Import node as ``resolved_module`` so
        consumers of the graph can recover the importing statement's line.
        """
        resolver = _ImportResolver(
            {
                key[len("module:") :]
                for key in self._key_to_id
                if key.startswith("module:")
            }
        )
        linked: set[tuple[int, int]] = set()
        for (
            module_id,
            import_id,
            module_name,
            imported_name,
            level,
        ) in self._pending_imports:
            importer = self._nodes[module_id].path
            target = resolver.resolve(importer, module_name, imported_name, level)
            if target is None or target == importer:
                continue
            self._nodes[import_id].properties["resolved_module"] = target
            target_id = self._key_to_id[f"module:{target}"]
            if (module_id, target_id) not in linked:
                linked.add((module_id, target_id))
                self._add_edge(module_id, target_id, CIGEdgeType.IMPORTS)

    # --- construction ---

    def replay(self, ops: list[tuple]) -> None:
        """Apply one module's op log, resolving edges against current state.

        Replays must run in forest order so each module's deferred edge
        resolution sees exactly the keys the serial build would have seen.
        """
        local_to_global: list[int] = []
        for op in ops:
            kind = op[0]
            if kind == "n":
                local_to_global.append(
                    self._add_node(op[1], op[2], op[3], op[4], op[5], op[6], op[7])
                )
            elif kind == "e":
                self._add_edge(local_to_global[op[1]], local_to_global[op[2]], op[3])
            elif kind == "ek":
                self.sink_edge_to_key(local_to_global[op[1]], op[2], op[3])
            elif kind == "ke":
                self.sink_edge_from_key(op[1], local_to_global[op[2]], op[3])
            elif kind == "c":
                self.sink_call_edge(local_to_global[op[1]], op[2])
            elif kind == "i":
                self.sink_import(
                    local_to_global[op[1]], local_to_global[op[2]], *op[3:]
                )

    def build(self) -> CodeIntelligenceGraph:
        """Construct the CIG with CSR index from accumulated nodes and edges."""
        self._resolve_imports()
        node_count = self._next_id
        csr = _CSRIndex(node_count, self._edges_list)

        return CodeIntelligenceGraph(
            csr=csr,
            nodes=self._nodes,
            key_to_id=self._key_to_id,
            edges_by_type=self._edges_by_type,
        )


class _ModuleRecorder:
    """Sink that records a module's nodes/edges as an op log (thread-local).

    Cross-module edge resolution is deferred as op records and decided
    during serial replay, so worker threads need no shared state.
    """

    def __init__(self) -> None:
        self.ops: list[tuple] = []
        self._next_local = 0

    def sink_node(
        self,
        key: str,
        node_type: str,
        name: str,
        qualified_name: str,
        path: str,
        line: int,
        properties: dict[str, Any] | None = None,
    ) -> int:
        local_id = self._next_local
        self._next_local += 1
        self.ops.append(
            ("n", key, node_type, name, qualified_name, path, line, properties)
        )
        return local_id

    def sink_edge(self, source_id: int, target_id: int, edge_type: CIGEdgeType) -> None:
        self.ops.append(("e", source_id, target_id, edge_type))

    def sink_edge_to_key(
        self, source_id: int, target_key: str, edge_type: CIGEdgeType
    ) -> None:
        self.ops.append(("ek", source_id, target_key, edge_type))

    def sink_edge_from_key(
        self, target_key: str, dest_id: int, edge_type: CIGEdgeType
    ) -> None:
        self.ops.append(("ke", target_key, dest_id, edge_type))

    def sink_call_edge(self, source_id: int, callee_name: str) -> None:
        self.ops.append(("c", source_id, callee_name))

    def sink_import(
        self,
        module_id: int,
        import_id: int,
        module_name: str,
        imported_name: str | None,
        level: int,
    ) -> None:
        self.ops.append(("i", module_id, import_id, module_name, imported_name, level))


class _AstExtractor:
    """Mirror of _cig_builder._ModuleVisitor over CPython ast trees."""

    def __init__(self, module_id: int, path: str, sink: _CIGSink) -> None:
        self._module_id = module_id
        self._path = path
        self._sink = sink
        self._class_stack: list[int] = []
        self._current_func_id: int | None = None

    # --- traversal ---

    def walk(self, root: ast.AST) -> None:
        """Iterative pre-order walk in CST-matching child order.

        Decorators and signature children are visited before the body for
        ClassDef/FunctionDef, matching libcst's source-order traversal so
        call attribution and node insertion order stay identical.
        """
        stack: list[tuple[ast.AST, bool]] = [(root, False)]
        while stack:
            node, leaving = stack.pop()
            if leaving:
                if isinstance(node, ast.ClassDef):
                    self._class_stack.pop() if self._class_stack else None
                elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    self._current_func_id = None
                continue
            if isinstance(node, ast.ClassDef):
                self._visit_class(node)
                stack.append((node, True))
                stack.extend((c, False) for c in self._class_children(node))
            elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                self._visit_func(node)
                stack.append((node, True))
                stack.extend((c, False) for c in self._func_children(node))
            else:
                if isinstance(node, ast.Import):
                    self._visit_import(node)
                elif isinstance(node, ast.ImportFrom):
                    self._visit_import_from(node)
                elif isinstance(node, ast.Call):
                    self._visit_call(node)
                stack.extend(
                    (c, False) for c in reversed(list(ast.iter_child_nodes(node)))
                )

    @staticmethod
    def _class_children(node: ast.ClassDef) -> list[ast.AST]:
        children: list[ast.AST] = list(node.decorator_list)
        children.extend(getattr(node, "type_params", []))
        children.extend(node.bases)
        children.extend(node.keywords)
        children.extend(node.body)
        return list(reversed(children))

    @staticmethod
    def _func_children(node: ast.FunctionDef | ast.AsyncFunctionDef) -> list[ast.AST]:
        children: list[ast.AST] = list(node.decorator_list)
        children.extend(getattr(node, "type_params", []))
        children.append(node.args)
        if node.returns is not None:
            children.append(node.returns)
        children.extend(node.body)
        return list(reversed(children))

    # --- ClassDef ---

    def _visit_class(self, node: ast.ClassDef) -> None:
        class_name = node.name
        qualified_name = f"{self._path}:{class_name}"
        key = f"class:{qualified_name}"

        is_protocol = False
        is_abstract = False
        is_dataclass = False
        bases: list[str] = []

        for base in node.bases:
            base_name = _extract_name(base)
            if base_name:
                bases.append(base_name)
                if base_name in ("Protocol", "typing.Protocol"):
                    is_protocol = True
                if base_name in ("ABC", "abc.ABC"):
                    is_abstract = True

        for dec in node.decorator_list:
            dec_name = _extract_name(dec)
            if dec_name == "dataclass":
                is_dataclass = True

        class_id = self._sink.sink_node(
            key=key,
            node_type="Class",
            name=class_name,
            qualified_name=qualified_name,
            path=self._path,
            line=node.lineno,
            properties={
                "is_protocol": is_protocol,
                "is_abstract": is_abstract,
                "is_dataclass": is_dataclass,
                "bases": bases,
            },
        )

        parent_id = self._class_stack[-1] if self._class_stack else self._module_id
        self._sink.sink_edge(parent_id, class_id, CIGEdgeType.DEFINES)
        self._class_stack.append(class_id)

    # --- FunctionDef / AsyncFunctionDef ---

    def _visit_func(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
        func_name = node.name
        qualified_name = f"{self._path}:{func_name}"

        is_method = bool(self._class_stack)
        node_type = "Method" if is_method else "Function"
        key = f"{node_type.lower()}:{qualified_name}"

        is_async = isinstance(node, ast.AsyncFunctionDef)
        return_type = _extract_annotation(node.returns)

        is_classmethod = False
        is_staticmethod = False
        is_property = False
        is_abstract = False
        for dec in node.decorator_list:
            dec_name = _extract_name(dec)
            if dec_name == "classmethod":
                is_classmethod = True
            elif dec_name == "staticmethod":
                is_staticmethod = True
            elif dec_name == "property":
                is_property = True
            elif dec_name == "abstractmethod":
                is_abstract = True

        if (
            is_method
            and self._path.startswith("tests/")
            or self._path.startswith("test_")
        ):
            if func_name.startswith("test_"):
                node_type = "Test"
                key = f"test:{qualified_name}"

        func_id = self._sink.sink_node(
            key=key,
            node_type=node_type,
            name=func_name,
            qualified_name=qualified_name,
            path=self._path,
            line=node.lineno,
            properties={
                "is_async": is_async,
                "return_type": return_type,
                "is_classmethod": is_classmethod,
                "is_staticmethod": is_staticmethod,
                "is_property": is_property,
                "is_abstract": is_abstract,
                "param_count": len(node.args.args),
            },
        )

        if is_method and self._class_stack:
            parent_id = self._class_stack[-1]
        else:
            parent_id = self._module_id
        self._sink.sink_edge(parent_id, func_id, CIGEdgeType.DEFINES)

        if node_type == "Test":
            target_name = func_name[5:] if func_name.startswith("test_") else func_name
            target_key = f"function:{self._path}:{target_name}"
            self._sink.sink_edge_from_key(target_key, func_id, CIGEdgeType.TESTED_BY)

        self._current_func_id = func_id

        # Positional-or-keyword params only (mirrors cst Parameters.params).
        # ast aligns defaults to the tail of posonlyargs + args.
        args = node.args
        combined = list(args.posonlyargs) + list(args.args)
        default_offset = len(combined) - len(args.defaults)
        for i, param in enumerate(args.args):
            default_expr = None
            k = len(args.posonlyargs) + i
            if k >= default_offset:
                default_expr = args.defaults[k - default_offset]
            param_default = _extract_name(default_expr) if default_expr else None

            param_key = f"param:{qualified_name}:{param.arg}"
            param_id = self._sink.sink_node(
                key=param_key,
                node_type="Parameter",
                name=param.arg,
                qualified_name=f"{qualified_name}.{param.arg}",
                path=self._path,
                line=node.lineno,
                properties={
                    "type_annotation": _extract_annotation(param.annotation),
                    "default": param_default,
                    "position": i,
                },
            )
            self._sink.sink_edge(func_id, param_id, CIGEdgeType.HAS_PARAMETER)

    # --- imports ---

    def _visit_import(self, node: ast.Import) -> None:
        for alias in node.names:
            self._record_import(alias.name, None, 0, node.lineno)

    def _visit_import_from(self, node: ast.ImportFrom) -> None:
        module_name = node.module or ""
        for alias in node.names:
            if alias.name == "*":
                continue
            self._record_import(module_name, alias.name, node.level, node.lineno)

    def _record_import(
        self, module_name: str, imported_name: str | None, level: int, line: int
    ) -> None:
        dotted = "." * level + module_name
        import_id = self._sink.sink_node(
            key=f"import:{self._path}:{dotted}:{imported_name}",
            node_type="Import",
            name=imported_name or dotted,
            qualified_name=(
                f"{dotted}.{imported_name}"
                if imported_name and module_name
                else f"{dotted}{imported_name or ''}"
            ),
            path=self._path,
            line=line,
            properties={
                "imported_module": module_name,
                "imported_name": imported_name,
                "is_relative": level > 0,
                "level": level,
            },
        )
        self._sink.sink_edge(self._module_id, import_id, CIGEdgeType.DEFINES)
        # The IMPORTS edge is resolved after every module is registered.
        self._sink.sink_import(
            self._module_id, import_id, module_name, imported_name, level
        )

    # --- calls ---

    def _visit_call(self, node: ast.Call) -> None:
        if self._current_func_id is None:
            return
        callee_name = _extract_call_name(node.func)
        if callee_name:
            self._sink.sink_call_edge(self._current_func_id, callee_name)


def _extract_name(node: ast.AST | None) -> str | None:
    """Leftmost simple name (mirrors _cig_builder._extract_name).

    Booleans and None are Name nodes in LibCST but Constants in ast;
    they are the only literals cst's extractor resolves, so they are the
    only Constants resolved here.
    """
    if node is None:
        return None
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Constant) and (
        isinstance(node.value, bool) or node.value is None
    ):
        return str(node.value)
    if isinstance(node, ast.Attribute):
        return _extract_name(node.value)
    if isinstance(node, ast.Call):
        return _extract_name(node.func)
    return None


def _extract_dotted_name(node: ast.AST | None) -> str | None:
    """Dotted name, e.g. 'foo.bar.baz' (None for anything else)."""
    if node is None:
        return None
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        prefix = _extract_dotted_name(node.value)
        if prefix:
            return f"{prefix}.{node.attr}"
    return None


def _extract_annotation(node: ast.AST | None) -> str | None:
    """Type annotation as string.

    Dotted name for simple annotations; ast.unparse for complex ones
    (the CST extractor stores the CST node repr there — the one
    intentional semantic delta of this backend).
    """
    if node is None:
        return None
    try:
        return _extract_dotted_name(node) or ast.unparse(node)
    except Exception:
        return None


def _extract_call_name(node: ast.AST) -> str | None:
    """Name of the called function (rightmost for attributes)."""
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        return node.attr
    return None


class _ImportResolver:
    """Map an import to the registered module it names.

    Absolute names (standard-library names excluded) are tried against, in
    order, the roots of the importer. A script or test directory without an
    ``__init__.py`` uses each ancestor directory (nearest first, ending at the
    repo root), each top-level directory (``src/``, ``scripts/`` ...) and each
    inferred package root (the parent of the outermost ``__init__.py`` chain).
    A module inside a package nested below the repo root (``src/pkg``) uses
    only its own package root, the repo root and the other package roots. Relative names resolve by level from the
    importing module's own package and never climb past the top package. A
    dotted name resolves to ``<root>/<dotted>.py`` or
    ``<root>/<dotted>/__init__.py``; the first root wins.
    """

    def __init__(self, paths: set[str]) -> None:
        self._paths = paths
        self._top_dirs = sorted({p.split("/", 1)[0] for p in paths if "/" in p})
        self._pkg_roots = sorted(
            {self._package_root(p) for p in paths if p.endswith("__init__.py")}
        )

    def _package_root(self, init_path: str) -> str:
        """Directory containing the outermost package above ``init_path``."""
        parts = init_path.split("/")[:-1]
        while len(parts) > 1 and "/".join([*parts[:-1], "__init__.py"]) in self._paths:
            parts = parts[:-1]
        return "/".join(parts[:-1])

    def _lookup(self, root: str, dotted: list[str]) -> str | None:
        base = "/".join(p for p in (root, *dotted) if p)
        for candidate in (f"{base}.py", f"{base}/__init__.py"):
            if candidate in self._paths:
                return candidate
        return None

    def resolve(
        self, importer: str, module_name: str, imported_name: str | None, level: int
    ) -> str | None:
        """Path of the module an import names, or None (external or unknown).

        ``from pkg import name`` prefers the submodule ``pkg.name`` and falls
        back to ``pkg`` itself (``name`` is then a symbol inside it).
        """
        parts = module_name.split(".") if module_name else []
        if not level and parts and parts[0] in sys.stdlib_module_names:
            return None
        names = [parts]
        if imported_name:
            names.insert(0, [*parts, imported_name])
        roots = (
            self._relative_roots(importer, level)
            if level
            else self._roots_for(importer)
        )
        # Roots outer: an earlier root's package beats a later root's submodule.
        for root in roots:
            for dotted in names:
                # ``dotted`` is empty only for ``from . import x`` (level > 0).
                found = self._lookup(root, dotted) if dotted else self._package(root)
                if found is not None:
                    return found
        return None

    def _package(self, root: str) -> str | None:
        init = f"{root}/__init__.py" if root else "__init__.py"
        return init if init in self._paths else None

    def _relative_roots(self, importer: str, level: int) -> list[str]:
        """Package directory a ``level``-dot relative import starts from."""
        package = importer.split("/")[:-1]
        keep = len(package) - (level - 1)
        if keep < 0:
            return []
        root = "/".join(package[:keep])
        # Climbing above the importing package must land in another package
        # (Python: "beyond top-level package" otherwise).
        if level > 1 and f"{root}/__init__.py" not in self._paths:
            return []
        return [root]

    def _roots_for(self, importer: str) -> list[str]:
        package = importer.split("/")[:-1]
        own_dir = "/".join(package)
        # A script or test dir (no __init__.py) is on sys.path itself, so its
        # ancestors are import roots; a package dir is not (no implicit
        # same-directory absolute imports, which would shadow the stdlib).
        in_package = bool(package) and f"{own_dir}/__init__.py" in self._paths
        own_root = self._package_root(f"{own_dir}/__init__.py") if in_package else ""
        if own_root:
            # A module inside a package nested below the repo root (``src/pkg``)
            # imports from its own package root, the repo root and the other
            # inferred package roots (installed packages): a bare name never
            # binds into a plain ``scripts/`` or ``tests/`` directory. A
            # repo-root package (``scripts/`` with an ``__init__.py``) is run
            # as a directory of scripts and keeps the top-level directories.
            return list(dict.fromkeys([own_root, "", *self._pkg_roots]))
        ancestors = (
            [""]
            if in_package
            else ["/".join(package[:n]) for n in range(len(package), -1, -1)]
        )
        ordered: dict[str, None] = {}
        for root in [*ancestors, *self._top_dirs, *self._pkg_roots]:
            ordered.setdefault(root)
        return list(ordered)