Spaces:
Sleeping
Sleeping
File size: 5,648 Bytes
b82e7a7 3d19c46 | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | import ast
def parse_functions_from_files(file_dict):
functions = {}
defined_funcs = set()
def infer_type_from_value(value_node):
if isinstance(value_node, ast.Call) and isinstance(value_node.func, ast.Attribute):
if value_node.func.attr in ("read_csv", "DataFrame"): return "pd.DataFrame"
if value_node.func.attr == "array": return "np.ndarray"
elif isinstance(value_node, ast.List): return "list"
elif isinstance(value_node, ast.Dict): return "dict"
elif isinstance(value_node, ast.Set): return "set"
elif isinstance(value_node, ast.Constant):
if isinstance(value_node.value, str): return "str"
if isinstance(value_node.value, bool): return "bool"
if isinstance(value_node.value, int): return "int"
if isinstance(value_node.value, float): return "float"
return "?"
for fname, code in file_dict.items():
tree = ast.parse(code)
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef): defined_funcs.add(node.name)
for fname, code in file_dict.items():
tree = ast.parse(code)
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef):
func_name = node.name
args, returns, calls = [], [], []
local_assignments = {}
reads_state, writes_state = set(), set()
for arg in node.args.args:
arg_type = ast.unparse(arg.annotation) if arg.annotation else "?"
args.append(f"{arg.arg}: {arg_type}")
for sub in ast.walk(node):
if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name):
if sub.func.id in defined_funcs: calls.append(sub.func.id)
elif isinstance(sub, ast.Assign):
for target in sub.targets:
if isinstance(target, ast.Name):
local_assignments[target.id] = infer_type_from_value(sub.value)
elif isinstance(sub, ast.Return):
if sub.value is None: continue
if isinstance(sub.value, ast.Tuple):
for elt in sub.value.elts:
label = ast.unparse(elt)
returns.append(f"{label}: {local_assignments.get(label, infer_type_from_value(elt))}")
else:
label = ast.unparse(sub.value)
returns.append(f"{label}: {local_assignments.get(label, infer_type_from_value(sub.value))}")
functions[func_name] = {
"args": args,
"returns": returns,
"calls": calls,
"filename": fname,
"reads_state": sorted(reads_state),
"writes_state": sorted(writes_state)
}
return functions
def get_reachable_functions(start, graph):
visited, stack = set(), [start]
while stack:
node = stack.pop()
if node not in visited:
visited.add(node)
stack.extend(graph.get(node, []))
return visited
def get_backtrace_functions(target, graph):
reverse_graph = {}
for caller, callees in graph.items():
for callee in callees:
reverse_graph.setdefault(callee, []).append(caller)
visited, stack = set(), [target]
while stack:
node = stack.pop()
if node not in visited:
visited.add(node)
stack.extend(reverse_graph.get(node, []))
return visited
from collections import deque, defaultdict
def build_nodes_and_edges(parsed, root_func, reachable, reverse=False, max_depth=10):
depth_map = {}
x_offset_map = defaultdict(int)
positions = {}
visited = set()
queue = deque([(root_func, 0)])
while queue:
current, depth = queue.popleft()
if current in visited or depth > max_depth:
continue
visited.add(current)
adjusted_depth = -depth if reverse else depth
x = x_offset_map[adjusted_depth] * 300
y = adjusted_depth * 150
positions[current] = {"x": x, "y": y}
x_offset_map[adjusted_depth] += 1
if reverse:
# Find callers of this function
next_funcs = [
caller for caller, meta in parsed.items()
if current in meta["calls"] and caller in reachable
]
else:
# Find callees
next_funcs = [
callee for callee in parsed[current]["calls"]
if callee in reachable
]
for nxt in next_funcs:
queue.append((nxt, depth + 1))
nodes = [{
"data": {"id": name, "label": name},
"position": positions.get(name, {"x": 0, "y": 0}),
"classes": "main" if name == root_func else ""
} for name in visited]
edges = []
for src in visited:
call_sequence = parsed[src]["calls"]
call_index = 1
for tgt in call_sequence:
if tgt in visited:
edges.append({
"data": {
"source": src,
"target": tgt,
"label": str(call_index)
},
"style": {
"line-width": 4 if call_sequence.count(tgt) > 1 else 2
}
})
call_index += 1
return nodes, edges
|