Coverage for src / renaissance / integrations / python / ast / cst_node.py: 93%

1from pathlib import Path 

2from typing import Self 

3 

4import libcst 

5from libcst import BaseCompoundStatement, BaseSmallStatement, ClassDef, CSTNode, FunctionDef, MetadataWrapper 

6from libcst.metadata import WhitespaceInclusivePositionProvider 

7 

8from renaissance.integrations.python.ast.util import convert 

9from renaissance.integrations.types import KIND_MAP, UnknownType 

10from renaissance.utils.ast_utils import next_sibling, preceding_sibling 

11 

12 

13class PythonCstTranslationUnit: 

14 def __init__(self, content, file_name: str): 

15 self.content = content 

16 self.lines = content.splitlines() 

17 self.file_name = file_name 

18 self.references_initialized = False 

19 self.wrapper = MetadataWrapper(libcst.parse_module(content)) 

20 self.atu = self.wrapper.module 

21 self.spans = self.wrapper.resolve(WhitespaceInclusivePositionProvider) 

22 

23 def start_of(self, node: CSTNode) -> int: 

24 span = self.spans.get(node) 

25 return convert(self.lines, span.start.line, span.start.column) if span else 0 

26 

27 def end_of(self, node: CSTNode) -> int: 

28 span = self.spans.get(node) 

29 return convert(self.lines, span.end.line, span.end.column) if span else 0 

30 

31 def signature_of(self, node: CSTNode) -> str: 

32 try: 

33 return self.atu.code_for_node(node) 

34 except Exception: 

35 return "" 

36 

37 

38class PythonCstNode: 

39 def __init__(self, node: CSTNode, translation_unit: PythonCstTranslationUnit, parent=None): 

40 self.parent = parent 

41 if parent and parent.root: 

42 self.root = parent.root 

43 else: 

44 self.root = self 

45 self.translation_unit = translation_unit 

46 self.node = node 

47 

48 self.is_statement = isinstance(self.node, (BaseSmallStatement, BaseCompoundStatement)) 

49 

50 # for matcher 

51 self.ast_type = KIND_MAP.get(type(node).__name__, UnknownType) # type(node)) 

52 if self.ast_type == UnknownType: 

53 print(f'"{type(node).__name__}": {type(node).__name__},') 

54 self.children: list[Self] = [PythonCstNode(node, translation_unit, self) for node in node.children] 

55 self.properties = {} 

56 

57 # for shower 

58 self.is_implicit = True 

59 self.show_props = False 

60 

61 # for rewriter 

62 self.text = self.signature 

63 

64 def __str__(self): 

65 return str(self.node) 

66 

67 def __repr__(self): 

68 return repr(self.node) 

69 

70 @property 

71 def signature(self): 

72 return self.translation_unit.signature_of(self.node) 

73 

74 @property 

75 def offset(self): 

76 return self.translation_unit.start_of(self.node) 

77 

78 @property 

79 def length(self): 

80 return self.end_offset - self.offset 

81 

82 @property 

83 def end_offset(self): 

84 return self.translation_unit.end_of(self.node) 

85 

86 @property 

87 def filename(self): 

88 return self.translation_unit.file_name 

89 

90 @property 

91 def name(self): 

92 if isinstance(self.node, (ClassDef, FunctionDef)): 

93 return self.node.name.value 

94 return "" 

95 self.name = "" # self._derive_name() 

96 

97 @property 

98 def next_sibling(self) -> Self | None: 

99 return next_sibling(self) 

100 

101 @property 

102 def preceding_sibling(self) -> Self | None: 

103 return preceding_sibling(self) 

104 

105 @staticmethod 

106 def load(file_path: Path) -> PythonCstNode: 

107 with Path(file_path).open() as file: 

108 content = file.read() 

109 return PythonCstNode.load_from_text(content, str(file_path)) 

110 

111 @staticmethod 

112 def load_from_text( 

113 text: str, 

114 file_name: str = "cst_snippet.py", 

115 ) -> PythonCstNode: 

116 translation_unit = PythonCstTranslationUnit(text, file_name=str(file_name)) 

117 root_node = PythonCstNode(translation_unit.atu, translation_unit) 

118 return root_node