Coverage for src/renaissance/integrations/python/ast/extractor.py: 84%

1"""AI: Extractor that builds a codebase graph of Python modules and their contained statements.""" 

2 

3from pathlib import Path 

4 

5import networkx as nx 

6 

7from renaissance.integrations.python.ast.rst_node import PythonRstNode 

8from renaissance.syntax_tree.semantic_kind import SemanticKind 

9 

10 

11class PythonExtractor: 

12 """AI: Extractor that builds a codebase graph of Python modules and their contained statements.""" 

13 

14 graph = nx.DiGraph() 

15 codebase: dict = {} 

16 

17 def process(self, file: Path): 

18 """AI: Parse file and merge its module/import/function/class nodes into the codebase graph.""" 

19 root = PythonRstNode.load(file) 

20 module_name = root.filename.replace("/", ".").replace(".py", "") 

21 folder = str(Path(file).parent) 

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

23 self.graph.add_edge(folder, module_name, type="contains") 

24 

25 for stmt in root: 

26 match stmt.semantic_kind: 

27 case SemanticKind.IMPORT: 

28 self.graph.add_edge(module_name, stmt.name, type="include") 

29 case _ if stmt.parser_kind == "ImportFrom": 

30 for alias in stmt.node.names: 

31 self.graph.add_edge(module_name, f"{stmt.node.module}.{alias.name}", type="include") 

32 case SemanticKind.FUNCTION: 

33 self.graph.add_edge(module_name, f"{module_name}.{stmt.name}", type="definition") 

34 self.graph.add_node(f"{module_name}.{stmt.name}", properties="function") 

35 # TODO: convert #, stmt.properties) to graphml 

36 case SemanticKind.CLASS: 

37 self.graph.add_edge(module_name, f"{module_name}.{stmt.name}", type="definition") 

38 self.graph.add_node(f"{module_name}.{stmt.name}") # convert to args, stmt.properties) 

39 case _: 

40 pass 

41 

42 self.codebase[file] = root 

43 # # reconstruct dependencies inside module 

44 # tu.lazy_create_refers(root) 

45 # self.nodes |= tu._nodes 

46 # self.edges |=tu._references 

47 # self.edges |= tu._referenced_by 

48 

49 def save_graph(self, filename: str): 

50 """AI: Write the extracted codebase graph to a GraphML file named filename.""" 

51 nx.write_graphml(self.graph, filename) 

52 print(f"Graph saved to: {filename}")