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

1from pathlib import Path 

2 

3import networkx 

4 

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

6 

7 

8class PythonExtractor: 

9 graph = networkx.DiGraph() 

10 codebase: dict = {} 

11 

12 def process(self, file: Path): 

13 root = PythonRstNode.load(file) 

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

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

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

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

18 

19 for stmt in root: 

20 match stmt.ast_type: 

21 case "Import": 

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

23 case "ImportFrom": 

24 for alias in stmt.node.names: 

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

26 case "FunctionDef": 

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

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

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

30 case "ClassDef": 

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

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

33 case _: 

34 pass 

35 

36 self.codebase[file] = root 

37 # # reconstruct dependencies inside module 

38 # tu.lazy_create_refers(root) 

39 # self.nodes |= tu._nodes 

40 # self.edges |=tu._references 

41 # self.edges |= tu._referenced_by 

42 

43 def save_graph(self, filename: str): 

44 networkx.write_graphml(self.graph, filename) 

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