Spaces:
Sleeping
Sleeping
File size: 5,959 Bytes
b82e7a7 3145636 b82e7a7 3145636 b82e7a7 3145636 b82e7a7 3145636 b82e7a7 3145636 b82e7a7 3145636 b82e7a7 3145636 b82e7a7 3145636 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 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | import ast
import ast
class CallCollector(ast.NodeVisitor):
def __init__(self, defined_funcs):
self.calls = []
self.defined_funcs = defined_funcs
def visit_Call(self, node):
if isinstance(node.func, ast.Name) and node.func.id in self.defined_funcs:
self.calls.append(node.func.id)
self.generic_visit(node)
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 = [], []
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}")
collector = CallCollector(defined_funcs)
collector.visit(node)
calls = collector.calls
for sub in ast.walk(node):
if 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
|