Coverage for skema/program_analysis/CAST/matlab/node_helper.py: 100%

1from typing import List, Dict 

2from skema.program_analysis.CAST2FN.model.cast import SourceRef 

3from skema.program_analysis.CAST.matlab.tokens import OTHER_TOKENS, KEYWORDS 

4 

5from tree_sitter import Node 

6 

7CONTROL_CHARACTERS = OTHER_TOKENS 

8 

9class NodeHelper(): 

10 def __init__(self, source: str, source_file_name: str): 

11 self.source = source 

12 self.source_file_name = source_file_name 

13 

14 def get_source_ref(self, node: Node) -> SourceRef: 

15 """Given a node and file name, return a CAST SourceRef object.""" 

16 row_start, col_start = node.start_point 

17 row_end, col_end = node.end_point 

18 return SourceRef(self.source_file_name, col_start, col_end, row_start, row_end) 

19 

20 def get_identifier(self, node: Node) -> str: 

21 """Given a node, return the identifier it represents. ie. The code between node.start_point and node.end_point""" 

22 line_num = 0 

23 column_num = 0 

24 in_identifier = False 

25 identifier = "" 

26 for i, char in enumerate(self.source): 

27 if line_num == node.start_point[0] and column_num == node.start_point[1]: 

28 in_identifier = True 

29 elif line_num == node.end_point[0] and column_num == node.end_point[1]: 

30 break 

31 

32 if char == "\n": 

33 line_num += 1 

34 column_num = 0 

35 else: 

36 column_num += 1 

37 

38 if in_identifier: 

39 identifier += char 

40 

41 return identifier 

42 

43def get_first_child_by_type(node: Node, child_type: str): 

44 """ Return the first node child with matching type """ 

45 for child in node.children: 

46 if child.type == child_type: 

47 return child 

48 

49 return None 

50 

51def get_children_by_types(node: Node, types: List): 

52 """Takes in a node and a list of types as inputs and returns all children matching those types. Otherwise, return an empty list""" 

53 return [child for child in node.children if child.type in types] 

54 

55def get_control_children(node: Node): 

56 """ return node children with control character types """ 

57 return get_children_by_types(node, CONTROL_CHARACTERS) 

58 

59def get_keyword_children(node: Node): 

60 """ return node children with Tree-sitter syntax keyword types """ 

61 return get_children_by_types(node, KEYWORDS)