Coverage for call_graph/parsers/python.py: 90%

1"""Python call graph parser."""

2

3 from __future__ import annotations

4

5 from tree_sitter import Node

6

7 from call_graph.parsers.base import CallGraphParser

8 from models import FunctionDefinition, CallEdge

9 from patterns.common import create_query, run_captures

10 from parsers.python import PY_LANGUAGE

11

12

13 class PythonCallParser(CallGraphParser):

14 def extract_function_definitions(self, root, source, filepath):

15 query_str = """(function_definition name: (identifier) @func_name)

16 (function_definition parameters: (parameters (identifier) @param))

17 (class_definition name: (identifier) @class_name)"""

18

19 query = create_query(PY_LANGUAGE, query_str)

20 captures = run_captures(query, root)

21

22 names = []

23 params = []

24 classes = []

25 name_to_params = {}

26

27 for node, capture_name in captures:

28 text = node.text.decode() if node.text else ""

29 line = node.start_point[0] + 1

30 if capture_name == "class_name":

31 classes.append((text, line))

32 elif capture_name == "func_name":

33 names.append((text, line))

34 name_to_params[(text, line)] = []

35 elif capture_name == "param":

36 if names:

37 names[-1]

38 params.append((text, line))

39

40 source_lines = source.split("\n")

41

42 for param, param_line in params:

43 best_name = None

44 min_gap = float("inf")

45

46 for name, name_line in names:

47 if name_line <= param_line <= name_line + 5:

48 gap = param_line - name_line

49 if gap < min_gap:

50 min_gap = gap

51 best_name = (name, name_line)

52

53 if best_name and best_name in name_to_params:

54 name_to_params[best_name].append(param)

55

56 defs = []

57 for name, name_line in names:

58 def_type = "function"

59 class_name = None

60

61 for cls, cls_line in classes:

62 if cls_line == name_line - 1:

63 def_type = "method"

64 class_name = cls

65 break

66

67 func_def = FunctionDefinition(

68 name=name,

69 filepath=str(filepath),

70 line=name_line,

71 def_type=def_type,

72 class_name=class_name,

73 params=name_to_params.get((name, name_line), []),

74 )

75 defs.append(func_def)

76

77 return defs

78

79 def extract_calls(self, root, source, filepath, all_functions):

80 source_lines = source.split("\n")

81 all_func_map = {(f.name, f.filepath): f for f in all_functions}

82

83 query_str = """(call) @call"""

84 query = create_query(PY_LANGUAGE, query_str)

85 captures = run_captures(query, root)

86

87 calls = []

88 for node, _ in captures:

89 call_site_line = node.start_point[0] + 1

90

91 caller_info = self._get_caller_info(

92 node, source_lines, all_func_map, filepath

93 )

94 if not caller_info:

95 continue

96

97 callee_name, receiver, is_method = self._get_callee_info(node)

98

99 call_type = "local"

100 if receiver:

101 call_type = "method"

102 elif callee_name not in [f.name for f in all_functions]:

103 call_type = "external"

104

105 call_edge = CallEdge(

106 caller_file=caller_info["filepath"],

107 caller_function=caller_info["name"],

108 caller_line=caller_info["line"],

109 callee_file=None,

110 callee_function=callee_name,

111 callee_line=None,

112 call_site_line=call_site_line,

113 call_type=call_type,

114 receiver_object=receiver,

115 )

116 calls.append(call_edge)

117

118 return calls

119

120 def _get_caller_info(self, call_node, source_lines, all_func_map, filepath):

121 line_num = call_node.start_point[0] + 1

122 filepath_str = str(filepath)

123

124 for func_key, func_def in all_func_map.items():

125 if func_def.filepath != filepath_str:

126 continue

127 if line_num >= func_def.line and line_num <= self._get_func_end_line(

128 func_def.line, source_lines

129 ):

130 return {

131 "name": func_def.name,

132 "filepath": func_def.filepath,

133 "line": func_def.line,

134 }

135

136 return None

137

138 def _get_func_end_line(self, start_line, source_lines):

139 for i in range(start_line, len(source_lines)):

140 line = source_lines[i]

141 if self._is_top_level_def(line):

142 return i

143 return len(source_lines)

144

145 def _is_top_level_def(self, line):

146 stripped = line.lstrip()

147 if stripped.startswith("def ") and not stripped.startswith(" "):

148 return True

149 if stripped.startswith("class ") and not stripped.startswith(" "):

150 return True

151 return False

152

153 def _get_callee_info(self, call_node):

154 for child in call_node.children:

155 if child.type == "identifier":

156 return child.text.decode(), None, False

157 if child.type == "attribute":

158 return self._get_attribute_callee(child)

159 if child.type == "member":

160 if child.parent.type == "call":

161 return self._get_member_callee(child)

162 return "<unknown>", None, False

163

164 def _get_member_callee(self, member_node):

165 member_text = member_node.text.decode()

166 return member_text, None, False

167

168 def _get_attribute_callee(self, attr_node):

169 parts = []

170

171 def extract_parts(node):

172 if node.type == "identifier":

173 parts.append(node.text.decode())

174 return

175 if node.type == "attribute":

176 for c in node.children:

177 if c.type != ".":

178 extract_parts(c)

179

180 extract_parts(attr_node)

181

182 if not parts:

183 return "<unknown>", None, False

184 if len(parts) >= 2:

185 receiver = ".".join(parts[:-1])

186 return parts[-1], receiver, True

187 return parts[0], None, False