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

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 of Python modules and their contained statements.""" 

8 

9from pathlib import Path 

10 

11import networkx as nx 

12 

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

14from renaissance.syntax_tree.semantic_kind import SemanticKind 

15 

16 

17class PythonExtractor: 

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

19 

20 graph = nx.DiGraph() 

21 codebase: dict = {} 

22 

23 def process(self, file: Path) -> None: 

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

25 root = PythonRstNode.load(file) 

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

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

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

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

30 

31 for stmt in root: 

32 match stmt.semantic_kind: 

33 case SemanticKind.IMPORT: 

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

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

36 for alias in stmt.node.names: 

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

38 case SemanticKind.FUNCTION: 

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

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

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

42 case SemanticKind.CLASS: 

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

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

45 case _: 

46 pass 

47 

48 self.codebase[file] = root 

49 # # reconstruct dependencies inside module 

50 # tu.lazy_create_refers(root) 

51 # self.nodes |= tu._nodes 

52 # self.edges |=tu._references 

53 # self.edges |= tu._referenced_by 

54 

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

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

57 nx.write_graphml(self.graph, filename) 

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