Spaces:
Sleeping
Sleeping
File size: 3,558 Bytes
46a9b7a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import ast
from dataclasses import dataclass
@dataclass
class Graph:
nodes: list[dict[str, str]]
edges: list[dict[str, str]]
def _call_target_name(node: ast.AST) -> str | None:
if isinstance(node, ast.Name):
return node.id
if isinstance(node, ast.Attribute) and isinstance(node.attr, str):
return node.attr
return None
def build_graph(code: str) -> Graph:
tree = ast.parse(code)
functions = {}
classes: dict[str, list[str]] = {}
for node in tree.body:
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
functions[node.name] = node
elif isinstance(node, ast.ClassDef):
method_names = []
for child in node.body:
if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
method_names.append(child.name)
classes[node.name] = method_names
nodes: list[dict[str, str]] = []
edges: list[dict[str, str]] = []
for class_name, method_names in classes.items():
nodes.append({"id": class_name, "label": class_name, "type": "class"})
for method_name in method_names:
method_id = f"{class_name}.{method_name}"
nodes.append({"id": method_id, "label": method_name, "type": "method"})
edges.append({"source": class_name, "target": method_id, "type": "contains"})
for func_name in functions:
nodes.append({"id": func_name, "label": func_name, "type": "function"})
known_nodes = {node["id"] for node in nodes}
call_edges = set()
def add_call_edge(source: str, target: str) -> None:
if source == target:
return
if target not in known_nodes:
return
call_edges.add((source, target))
for func_name, func_node in functions.items():
for call in [n for n in ast.walk(func_node) if isinstance(n, ast.Call)]:
target = _call_target_name(call.func)
if target is None:
continue
if target in functions:
add_call_edge(func_name, target)
for class_name, method_names in classes.items():
for node in tree.body:
if isinstance(node, ast.ClassDef) and node.name == class_name:
for child in node.body:
if not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
continue
source_id = f"{class_name}.{child.name}"
for call in [n for n in ast.walk(child) if isinstance(n, ast.Call)]:
target = call.func
if isinstance(target, ast.Name):
if target.id in functions:
add_call_edge(source_id, target.id)
continue
if isinstance(target, ast.Attribute):
if isinstance(target.value, ast.Name):
if target.value.id == "self":
target_id = f"{class_name}.{target.attr}"
add_call_edge(source_id, target_id)
elif target.value.id in classes:
target_id = f"{target.value.id}.{target.attr}"
add_call_edge(source_id, target_id)
for source, target in sorted(call_edges):
edges.append({"source": source, "target": target, "type": "calls"})
return Graph(nodes=nodes, edges=edges)
|