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
Dependency diagram
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
"""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)