Coverage for src/renaissance/integrations/tree_sitter/extractor.py: 100%

1"""AI: Extractor that builds a codebase graph from tree-sitter-matched patterns.""" 

2 

3from pathlib import Path 

4 

5import networkx as nx 

6 

7from renaissance.integrations.tree_sitter.adapter import TreeSitterAdapter 

8from renaissance.integrations.tree_sitter.factory import TreeSitterPatternFactory 

9from renaissance.syntax_tree import PatternMatch 

10from renaissance.syntax_tree.match_finder import match_pattern 

11from renaissance.syntax_tree.semantic_kind import SemanticKind 

12 

13GRAPHML_DIR = "out_graphml" 

14Path(GRAPHML_DIR).mkdir(parents=True, exist_ok=True) 

15 

16 

17def _has_semantic_kind(node, kind: SemanticKind) -> bool: 

18 return node.semantic_kind is kind 

19 

20 

21class Extractor: 

22 """AI: Find occurrences of a set of patterns in tree-sitter-parsed code.""" 

23 

24 def __init__(self, factory: TreeSitterPatternFactory, patterns: list[str]): 

25 """AI: Configure an extractor that finds occurrences of the given patterns in code.""" 

26 self.pattern_factory = factory 

27 self.patterns = patterns 

28 

29 def run(self, raw: str) -> list[PatternMatch]: 

30 """AI: Find occurrences of the configured patterns in raw code and return the matches.""" 

31 code = self.pattern_factory.create_statements(raw) 

32 results = [] 

33 for rule in self.patterns: 

34 pattern = self.pattern_factory.create_statements(rule) 

35 results.extend(match_pattern(code, pattern, {})) 

36 return results 

37 

38 

39class BaseCodeGraphExtractor: 

40 """AI: Base class that extracts a code graph from a language's tree-sitter-parsed files.""" 

41 

42 def __init__(self, language: str, lib_path: str): 

43 """AI: Configure a code-graph extractor for the given language and grammar library.""" 

44 self.language = language 

45 self.lib_path = lib_path 

46 self.adapter = TreeSitterAdapter(lib_path) 

47 self.graph = nx.DiGraph() 

48 

49 def extract(self, files): 

50 """AI: Parse each file and build the code graph by processing its LST.""" 

51 for f in files: 

52 try: 

53 code = Path(f).read_text() 

54 tree = self.adapter.parse_code(code) 

55 lst = self.adapter.to_lst(code, tree) 

56 self._process_file(f, lst) 

57 except Exception as e: 

58 print(f"Error processing {f}: {e}") 

59 

60 def _process_file(self, file_path: str, lst): 

61 raise NotImplementedError 

62 

63 def save_graph(self, filename: str): 

64 """AI: Write the extracted code graph to a GraphML file named filename.""" 

65 path = Path(GRAPHML_DIR) / filename 

66 nx.write_graphml(self.graph, path) 

67 print(f"Graph saved to: {path}") 

68 

69 

70class PythonCodeGraphExtractor(BaseCodeGraphExtractor): 

71 """AI: Extract a code graph (functions, calls) from Python source files.""" 

72 

73 def _process_file(self, file_path, lst): 

74 folder = str(Path(file_path).parent) 

75 self.graph.add_node(file_path, type="file", folder=folder) 

76 self.graph.add_node(folder, type="folder") 

77 self.graph.add_edge(folder, file_path, type="contains") 

78 

79 for node in lst.traverse(): 

80 if _has_semantic_kind(node, SemanticKind.FUNCTION): 

81 name = node.signature.split("(")[0].split()[-1] 

82 self.graph.add_node(name, type="function", file=file_path) 

83 self.graph.add_edge(file_path, name, type="defines") 

84 

85 elif _has_semantic_kind(node, SemanticKind.CALL): 

86 call_target = node.signature.strip().split("(")[0] 

87 self.graph.add_node(call_target, type="call_target") 

88 self.graph.add_edge(file_path, call_target, type="calls") 

89 

90 

91class JavaCodeGraphExtractor(BaseCodeGraphExtractor): 

92 """AI: Extract a code graph (methods, calls) from Java source files.""" 

93 

94 def _process_file(self, file_path, lst): 

95 folder = str(Path(file_path).parent) 

96 self.graph.add_node(file_path, type="file", folder=folder) 

97 self.graph.add_node(folder, type="folder") 

98 self.graph.add_edge(folder, file_path, type="contains") 

99 

100 for node in lst.traverse(): 

101 if _has_semantic_kind(node, SemanticKind.FUNCTION): 

102 name = node.properties.get("name", "method") 

103 self.graph.add_node(name, type="method", file=file_path) 

104 self.graph.add_edge(file_path, name, type="defines") 

105 

106 elif _has_semantic_kind(node, SemanticKind.CALL): 

107 target = node.signature.strip().split("(")[0] 

108 self.graph.add_node(target, type="method_target") 

109 self.graph.add_edge(file_path, target, type="calls") 

110 

111 

112class CppCodeGraphExtractor(BaseCodeGraphExtractor): 

113 """AI: Extract a code graph (functions, calls) from C++ source files.""" 

114 

115 def _process_file(self, file_path, lst): 

116 folder = str(Path(file_path).parent) 

117 self.graph.add_node(file_path, type="file", folder=folder) 

118 self.graph.add_node(folder, type="folder") 

119 self.graph.add_edge(folder, file_path, type="contains") 

120 

121 for node in lst.traverse(): 

122 if _has_semantic_kind(node, SemanticKind.FUNCTION): 

123 name = node.properties.get("name", "func") 

124 self.graph.add_node(name, type="function", file=file_path) 

125 self.graph.add_edge(file_path, name, type="defines") 

126 

127 elif _has_semantic_kind(node, SemanticKind.CALL): 

128 call_expr = node.signature.strip().split("(")[0] 

129 self.graph.add_node(call_expr, type="call_target") 

130 self.graph.add_edge(file_path, call_expr, type="calls")