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

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: Helpers for finding AST nodes matching a predicate or semantic kind.""" 

8 

9import re 

10from collections.abc import Callable, Iterator, Sequence 

11 

12from renaissance.utils.ast_utils import traverse 

13 

14from .ast_node import ASTNode 

15from .node_protocol import NodeProtocol 

16from .semantic_kind import SemanticKind 

17 

18 

19class ASTFinder: 

20 """AI: Static helpers for finding AST nodes matching a predicate or semantic kind.""" 

21 

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

23 

24 @staticmethod 

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

26 """AI: Return all descendant nodes of ast_node for which function returns a truthy value or child iterator.""" 

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

28 

29 # @staticmethod 

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

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

32 

33 @staticmethod 

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

35 """AI: Return all descendant nodes of ast_node whose kind matches the given kind pattern.""" 

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

37 

38 @staticmethod 

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

40 """AI: Return whether ast_node's kind matches the given kind pattern.""" 

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

42 # get kind of the ast_node with only word characters 

43 if ast_node is None: 

44 return False 

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

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

47 return pattern.fullmatch(ast_kind) is not None 

48 

49 @staticmethod 

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

51 result = function(ast_node) 

52 if isinstance(result, bool) and result: 

53 yield ast_node 

54 elif isinstance(result, Iterator): 

55 yield from result 

56 for child in ast_node.children: 

57 yield from ASTFinder.__find_all(child, function) 

58 

59 @staticmethod 

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

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

62 node_kind = ast_node.kind or "" 

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

64 

65 if pattern.fullmatch(ast_kind): 

66 yield ast_node 

67 for child in ast_node.children: 

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

69 yield from ASTFinder.__matches_kind(child, pattern) 

70 

71 

72def find_nodes(ast_node: NodeProtocol, predicate) -> Sequence[NodeProtocol]: 

73 """AI: Return all descendants (and ast_node itself) matching predicate, via a full traversal.""" 

74 return [node for node in traverse(ast_node) if predicate(node)] 

75 

76 

77def matches_node(ast_node: NodeProtocol, predicate) -> bool: 

78 """AI: Return True if ast_node itself satisfies predicate.""" 

79 return predicate(ast_node) 

80 

81 

82def find_semantic_kind(ast_node: NodeProtocol, kind: SemanticKind) -> Sequence[NodeProtocol]: 

83 """AI: Return all nodes under ast_node whose semantic kind matches the given kind.""" 

84 return find_nodes(ast_node, lambda node: node.semantic_kind is kind)