Coverage for src / renaissance / syntax_tree / ast_finder.py: 69%

1import re 

2from collections.abc import Callable, Iterator, Sequence 

3 

4from renaissance.integrations.types import Type 

5from renaissance.utils.ast_utils import traverse 

6 

7from .ast_node import ASTNode 

8 

9 

10class ASTFinder: 

11 KIND_MATCH = re.compile(r"[\W_]+") 

12 

13 @staticmethod 

14 def find_all(ast_node: ASTNode, function: Callable[[ASTNode], Iterator[ASTNode] | bool]) -> Sequence[ASTNode]: 

15 return list(ASTFinder.__find_all(ast_node, function)) 

16 

17 # @staticmethod 

18 # def find_kind(ast_node: ASTNode, kind: str | re.Pattern[str]) -> Sequence[ASTNode]: 

19 # return list(ASTFinder.__matches_kind(ast_node, kind)) 

20 

21 @staticmethod 

22 def find(ast_node: ASTNode, kind: str | re.Pattern[str]) -> Sequence[ASTNode]: 

23 return list(ASTFinder.__matches_kind(ast_node, kind)) 

24 

25 @staticmethod 

26 def matches_kind(ast_node: ASTNode | None, kind: str | re.Pattern[str]) -> bool: 

27 # compare kind with the ast_node kind only using word characters 

28 # get kind of the ast_node with only word characters 

29 if ast_node is None: 

30 return False 

31 ast_kind = ASTFinder.KIND_MATCH.sub("", ast_node.kind).lower() 

32 pattern = kind if isinstance(kind, re.Pattern) else re.compile(kind, re.IGNORECASE) 

33 return pattern.fullmatch(ast_kind) is not None 

34 

35 @staticmethod 

36 def __find_all(ast_node: ASTNode, function: Callable[[ASTNode], Iterator[ASTNode] | bool]) -> Iterator[ASTNode]: 

37 result = function(ast_node) 

38 if isinstance(result, bool) and result: 

39 yield ast_node 

40 elif isinstance(result, Iterator): 

41 yield from result 

42 for child in ast_node.children: 

43 yield from ASTFinder.__find_all(child, function) 

44 

45 @staticmethod 

46 def __matches_kind(ast_node: ASTNode, kind: str | re.Pattern[str]) -> Iterator[ASTNode]: 

47 pattern = kind if isinstance(kind, re.Pattern) else re.compile(kind, re.IGNORECASE) 

48 node_kind = ast_node.kind or "" 

49 ast_kind = ASTFinder.KIND_MATCH.sub("", node_kind).lower() 

50 

51 if pattern.fullmatch(ast_kind): 

52 yield ast_node 

53 for child in ast_node.children: 

54 # assert isinstance(child, type(ast_node)), f'Expected {type(ast_node)} but got {type(child)}' 

55 yield from ASTFinder.__matches_kind(child, pattern) 

56 

57 

58def find_ast_type(ast_node, kind: type[Type]) -> Sequence: 

59 return [n for n in traverse(ast_node) if isinstance(n.ast_type(), kind)] 

60 

61 

62def matches_kind(ast_node, kind: type[Type]) -> bool: 

63 return isinstance(ast_node.ast_type(), kind)