"""
Graph models for call graph representation.
These models are backend-agnostic and can be populated from either
Python's AST or LPython's ASR.
"""
from __future__ import annotations
from collections.abc import Iterator
from dataclasses import dataclass, field
from enum import StrEnum
[docs]
class NodeKind(StrEnum):
"""Kind of node in the call graph."""
MODULE = "module"
CLASS = "class"
FUNCTION = "function"
METHOD = "method"
LAMBDA = "lambda"
COMPREHENSION = "comprehension"
[docs]
class EdgeKind(StrEnum):
"""Kind of edge in the call graph."""
CALLS = "calls" # Direct function call
DEFINES = "defines" # Class/module defines function
INHERITS = "inherits" # Class inheritance
IMPORTS = "imports" # Module import
REFERENCES = "references" # Variable reference (weaker than call)
[docs]
@dataclass(frozen=True)
class SourceLocation:
"""Source code location."""
file: str
line_start: int
line_end: int
col_start: int = 0
col_end: int = 0
def __str__(self) -> str:
return f"{self.file}:{self.line_start}"
[docs]
def contains_line(self, line: int) -> bool:
"""Check if a line number falls within this location."""
return self.line_start <= line <= self.line_end
[docs]
@dataclass(frozen=True)
class FunctionSignature:
"""Function signature information."""
name: str
params: tuple[str, ...] = ()
return_type: str | None = None
decorators: tuple[str, ...] = ()
is_async: bool = False
is_generator: bool = False
@property
def arity(self) -> int:
"""Number of parameters."""
return len(self.params)
[docs]
@dataclass
class GraphNode:
"""Node in the call graph."""
id: str # Fully qualified name: module.class.function
kind: NodeKind
name: str # Short name
location: SourceLocation | None = None
signature: FunctionSignature | None = None
docstring: str | None = None
metadata: dict[str, any] = field(default_factory=dict)
def __hash__(self) -> int:
return hash(self.id)
def __eq__(self, other: object) -> bool:
if not isinstance(other, GraphNode):
return False
return self.id == other.id
[docs]
@dataclass(frozen=True)
class GraphEdge:
"""Edge in the call graph."""
source_id: str
target_id: str
kind: EdgeKind
location: SourceLocation | None = None # Where the call/reference occurs
is_conditional: bool = False # Inside if/try/loop
is_dynamic: bool = False # Dynamic call (getattr, etc.)
confidence: float = 1.0 # 1.0 = certain, <1.0 = inferred
def __hash__(self) -> int:
return hash((self.source_id, self.target_id, self.kind))
[docs]
@dataclass
class CallGraph:
"""
Directed graph representing function calls and definitions.
This is the core data structure for M2 graph-spectral analysis.
Designed to be populated from either Python AST or LPython ASR.
"""
nodes: dict[str, GraphNode] = field(default_factory=dict)
edges: set[GraphEdge] = field(default_factory=set)
# Index structures for efficient queries
_outgoing: dict[str, set[str]] = field(default_factory=dict)
_incoming: dict[str, set[str]] = field(default_factory=dict)
_by_file: dict[str, set[str]] = field(default_factory=dict)
[docs]
def add_node(self, node: GraphNode) -> None:
"""Add a node to the graph."""
self.nodes[node.id] = node
if node.id not in self._outgoing:
self._outgoing[node.id] = set()
if node.id not in self._incoming:
self._incoming[node.id] = set()
if node.location:
file_key = node.location.file
if file_key not in self._by_file:
self._by_file[file_key] = set()
self._by_file[file_key].add(node.id)
[docs]
def add_edge(self, edge: GraphEdge) -> None:
"""Add an edge to the graph."""
self.edges.add(edge)
if edge.source_id not in self._outgoing:
self._outgoing[edge.source_id] = set()
self._outgoing[edge.source_id].add(edge.target_id)
if edge.target_id not in self._incoming:
self._incoming[edge.target_id] = set()
self._incoming[edge.target_id].add(edge.source_id)
[docs]
def get_node(self, node_id: str) -> GraphNode | None:
"""Get a node by ID."""
return self.nodes.get(node_id)
[docs]
def get_callees(self, node_id: str) -> set[str]:
"""Get IDs of functions called by node_id."""
return self._outgoing.get(node_id, set())
[docs]
def get_callers(self, node_id: str) -> set[str]:
"""Get IDs of functions that call node_id."""
return self._incoming.get(node_id, set())
[docs]
def get_edges_from(self, node_id: str, kind: EdgeKind | None = None) -> Iterator[GraphEdge]:
"""Get all edges originating from a node."""
for edge in self.edges:
if edge.source_id == node_id:
if kind is None or edge.kind == kind:
yield edge
[docs]
def get_edges_to(self, node_id: str, kind: EdgeKind | None = None) -> Iterator[GraphEdge]:
"""Get all edges pointing to a node."""
for edge in self.edges:
if edge.target_id == node_id:
if kind is None or edge.kind == kind:
yield edge
[docs]
def get_nodes_in_file(self, file_path: str) -> set[str]:
"""Get all node IDs in a specific file."""
return self._by_file.get(file_path, set())
[docs]
def functions(self) -> Iterator[GraphNode]:
"""Iterate over all function/method nodes."""
for node in self.nodes.values():
if node.kind in (NodeKind.FUNCTION, NodeKind.METHOD, NodeKind.LAMBDA):
yield node
[docs]
def classes(self) -> Iterator[GraphNode]:
"""Iterate over all class nodes."""
for node in self.nodes.values():
if node.kind == NodeKind.CLASS:
yield node
[docs]
def modules(self) -> Iterator[GraphNode]:
"""Iterate over all module nodes."""
for node in self.nodes.values():
if node.kind == NodeKind.MODULE:
yield node
# ─────────────────────────────────────────────────────────────────
# Graph algorithms (M2 foundation)
# ─────────────────────────────────────────────────────────────────
[docs]
def reachable_from(self, node_id: str, max_depth: int | None = None) -> set[str]:
"""
Get all nodes reachable from node_id via calls.
Args:
node_id: Starting node
max_depth: Maximum traversal depth (None = unlimited)
Returns:
Set of reachable node IDs (excluding start node)
"""
visited: set[str] = set()
frontier = [(node_id, 0)]
while frontier:
current, depth = frontier.pop()
if current in visited:
continue
if current != node_id:
visited.add(current)
if max_depth is not None and depth >= max_depth:
continue
for callee in self.get_callees(current):
if callee not in visited:
frontier.append((callee, depth + 1))
return visited
[docs]
def reaches(self, node_id: str, max_depth: int | None = None) -> set[str]:
"""
Get all nodes that can reach node_id via calls.
Args:
node_id: Target node
max_depth: Maximum traversal depth (None = unlimited)
Returns:
Set of node IDs that can reach target (excluding target)
"""
visited: set[str] = set()
frontier = [(node_id, 0)]
while frontier:
current, depth = frontier.pop()
if current in visited:
continue
if current != node_id:
visited.add(current)
if max_depth is not None and depth >= max_depth:
continue
for caller in self.get_callers(current):
if caller not in visited:
frontier.append((caller, depth + 1))
return visited
[docs]
def strongly_connected_components(self) -> list[frozenset[str]]:
"""
Find strongly connected components using Tarjan's algorithm.
Returns:
List of SCCs, each as a frozenset of node IDs
"""
index_counter = [0]
stack: list[str] = []
lowlinks: dict[str, int] = {}
index: dict[str, int] = {}
on_stack: set[str] = set()
sccs: list[frozenset[str]] = []
def strongconnect(node: str) -> None:
index[node] = index_counter[0]
lowlinks[node] = index_counter[0]
index_counter[0] += 1
stack.append(node)
on_stack.add(node)
for callee in self.get_callees(node):
if callee not in index:
strongconnect(callee)
lowlinks[node] = min(lowlinks[node], lowlinks[callee])
elif callee in on_stack:
lowlinks[node] = min(lowlinks[node], index[callee])
if lowlinks[node] == index[node]:
scc: set[str] = set()
while True:
w = stack.pop()
on_stack.remove(w)
scc.add(w)
if w == node:
break
sccs.append(frozenset(scc))
for node in self.nodes:
if node not in index:
strongconnect(node)
return sccs
[docs]
def topological_sort(self) -> list[str]:
"""
Topological sort of nodes (only valid for DAG).
Returns:
List of node IDs in topological order
Raises:
ValueError: If graph contains cycles
"""
in_degree: dict[str, int] = dict.fromkeys(self.nodes, 0)
for edge in self.edges:
if edge.kind == EdgeKind.CALLS:
if edge.target_id in in_degree:
in_degree[edge.target_id] += 1
queue = [n for n, d in in_degree.items() if d == 0]
result: list[str] = []
while queue:
node = queue.pop(0)
result.append(node)
for callee in self.get_callees(node):
if callee in in_degree:
in_degree[callee] -= 1
if in_degree[callee] == 0:
queue.append(callee)
if len(result) != len(self.nodes):
raise ValueError("Graph contains cycles - not a DAG")
return result
[docs]
def condensation(self) -> "CallGraph":
"""
Create condensation graph (DAG of SCCs).
Each SCC becomes a single node in the condensation graph.
Useful for hierarchical decomposition (M2).
Returns:
New CallGraph where each node is an SCC
"""
sccs = self.strongly_connected_components()
node_to_scc: dict[str, int] = {}
for i, scc in enumerate(sccs):
for node in scc:
node_to_scc[node] = i
condensed = CallGraph()
# Create SCC nodes
for i, scc in enumerate(sccs):
scc_id = f"scc_{i}"
members = sorted(scc)
condensed.add_node(
GraphNode(
id=scc_id,
kind=NodeKind.MODULE, # SCC as pseudo-module
name=f"SCC({', '.join(members[:3])}{'...' if len(members) > 3 else ''})",
metadata={"members": list(scc), "size": len(scc)},
)
)
# Create edges between SCCs
seen_edges: set[tuple[int, int]] = set()
for edge in self.edges:
if edge.kind == EdgeKind.CALLS:
src_scc = node_to_scc.get(edge.source_id)
tgt_scc = node_to_scc.get(edge.target_id)
if src_scc is not None and tgt_scc is not None and src_scc != tgt_scc:
if (src_scc, tgt_scc) not in seen_edges:
seen_edges.add((src_scc, tgt_scc))
condensed.add_edge(
GraphEdge(
source_id=f"scc_{src_scc}",
target_id=f"scc_{tgt_scc}",
kind=EdgeKind.CALLS,
)
)
return condensed
# ─────────────────────────────────────────────────────────────────
# Serialization
# ─────────────────────────────────────────────────────────────────
[docs]
def to_dict(self) -> dict:
"""Serialize to dictionary."""
return {
"nodes": [
{
"id": n.id,
"kind": n.kind.value,
"name": n.name,
"location": {
"file": n.location.file,
"line_start": n.location.line_start,
"line_end": n.location.line_end,
}
if n.location
else None,
"signature": {
"name": n.signature.name,
"params": list(n.signature.params),
"return_type": n.signature.return_type,
"is_async": n.signature.is_async,
}
if n.signature
else None,
}
for n in self.nodes.values()
],
"edges": [
{
"source": e.source_id,
"target": e.target_id,
"kind": e.kind.value,
"confidence": e.confidence,
}
for e in self.edges
],
}
[docs]
@classmethod
def from_dict(cls, data: dict) -> "CallGraph":
"""Deserialize from dictionary."""
graph = cls()
for n in data.get("nodes", []):
location = None
if n.get("location"):
loc = n["location"]
location = SourceLocation(
file=loc["file"],
line_start=loc["line_start"],
line_end=loc["line_end"],
)
signature = None
if n.get("signature"):
sig = n["signature"]
signature = FunctionSignature(
name=sig["name"],
params=tuple(sig.get("params", [])),
return_type=sig.get("return_type"),
is_async=sig.get("is_async", False),
)
graph.add_node(
GraphNode(
id=n["id"],
kind=NodeKind(n["kind"]),
name=n["name"],
location=location,
signature=signature,
)
)
for e in data.get("edges", []):
graph.add_edge(
GraphEdge(
source_id=e["source"],
target_id=e["target"],
kind=EdgeKind(e["kind"]),
confidence=e.get("confidence", 1.0),
)
)
return graph
[docs]
def to_dot(self, title: str = "Call Graph") -> str:
"""Export to Graphviz DOT format."""
lines = [
f'digraph "{title}" {{',
" rankdir=TB;",
" node [shape=box, style=filled, fillcolor=lightblue];",
]
# Nodes
for node in self.nodes.values():
label = node.name
if node.signature:
params = ", ".join(node.signature.params[:3])
if len(node.signature.params) > 3:
params += "..."
label = f"{node.name}({params})"
color = {
NodeKind.MODULE: "lightgray",
NodeKind.CLASS: "lightyellow",
NodeKind.FUNCTION: "lightblue",
NodeKind.METHOD: "lightgreen",
NodeKind.LAMBDA: "pink",
}.get(node.kind, "white")
lines.append(f' "{node.id}" [label="{label}", fillcolor={color}];')
# Edges
for edge in self.edges:
style = "solid" if edge.confidence >= 0.9 else "dashed"
color = {
EdgeKind.CALLS: "black",
EdgeKind.DEFINES: "blue",
EdgeKind.INHERITS: "red",
EdgeKind.IMPORTS: "gray",
}.get(edge.kind, "black")
lines.append(f' "{edge.source_id}" -> "{edge.target_id}" [style={style}, color={color}];')
lines.append("}")
return "\n".join(lines)
def __len__(self) -> int:
return len(self.nodes)
def __repr__(self) -> str:
return f"CallGraph(nodes={len(self.nodes)}, edges={len(self.edges)})"