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

1# Copyright (C) 2024-2026 TNO 

2# 

3# Licensed under the terms specified in the LICENSE file in the root directory 

4# of this repository, also available at: 

5# https://github.com/TNO/Renaissance.Py/blob/main/LICENSE 

6 

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

8 

9from pathlib import Path 

10 

11import networkx as nx 

12 

13from renaissance.integrations.tree_sitter.adapter import TreeSitterAdapter 

14from renaissance.integrations.tree_sitter.factory import TreeSitterPatternFactory 

15from renaissance.syntax_tree import PatternMatch 

16from renaissance.syntax_tree.match_finder import match_pattern 

17from renaissance.syntax_tree.semantic_kind import SemanticKind 

18 

19GRAPHML_DIR = "out_graphml" 

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

21 

22 

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

24 return node.semantic_kind is kind 

25 

26 

27class Extractor: 

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

29 

30 def __init__(self, factory: TreeSitterPatternFactory, patterns: list[str]) -> None: 

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

32 self.pattern_factory = factory 

33 self.patterns = patterns 

34 

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

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

37 code = self.pattern_factory.create_statements(raw) 

38 results = [] 

39 for rule in self.patterns: 

40 pattern = self.pattern_factory.create_statements(rule) 

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

42 return results 

43 

44 

45class BaseCodeGraphExtractor: 

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

47 

48 def __init__(self, language: str, lib_path: str) -> None: 

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

50 self.language = language 

51 self.lib_path = lib_path 

52 self.adapter = TreeSitterAdapter(lib_path) 

53 self.graph = nx.DiGraph() 

54 

55 def extract(self, files) -> None: 

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

57 for f in files: 

58 try: 

59 code = Path(f).read_text() 

60 tree = self.adapter.parse_code(code) 

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

62 self._process_file(f, lst) 

63 except Exception as e: 

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

65 

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

67 raise NotImplementedError 

68 

69 def save_graph(self, filename: str) -> None: 

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

71 path = Path(GRAPHML_DIR) / filename 

72 nx.write_graphml(self.graph, path) 

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

74 

75 

76class PythonCodeGraphExtractor(BaseCodeGraphExtractor): 

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

78 

79 def _process_file(self, file_path, lst) -> None: 

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

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

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

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

84 

85 for node in lst.traverse(): 

86 if _has_semantic_kind(node, SemanticKind.FUNCTION): 

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

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

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

90 

91 elif _has_semantic_kind(node, SemanticKind.CALL): 

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

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

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

95 

96 

97class JavaCodeGraphExtractor(BaseCodeGraphExtractor): 

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

99 

100 def _process_file(self, file_path, lst) -> None: 

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

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

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

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

105 

106 for node in lst.traverse(): 

107 if _has_semantic_kind(node, SemanticKind.FUNCTION): 

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

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

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

111 

112 elif _has_semantic_kind(node, SemanticKind.CALL): 

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

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

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

116 

117 

118class CppCodeGraphExtractor(BaseCodeGraphExtractor): 

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

120 

121 def _process_file(self, file_path, lst) -> None: 

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

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

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

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

126 

127 for node in lst.traverse(): 

128 if _has_semantic_kind(node, SemanticKind.FUNCTION): 

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

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

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

132 

133 elif _has_semantic_kind(node, SemanticKind.CALL): 

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

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

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