Coverage for phml\utilities\travel\travel.py: 100%

37 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-04-12 14:26 -0500

1"""utilities.travel 

2 

3Collection of utilities that hep with traversing an ast or node tree. 

4""" 

5 

6from typing import Iterator 

7 

8from phml.nodes import Element, Node, Parent 

9 

10__all__ = ["path", "path_names", "walk", "visit_all_after"] 

11 

12 

13def path(node: Node) -> list[Element]: 

14 """Get a list of nodes where each one is a child of 

15 the other leading to the node passed in. This gives a 

16 path to the node. 

17 

18 Does not include given node. 

19 

20 Args: 

21 node (Node): Node to find ancestors of. 

22 

23 Returns: 

24 list[Node]: List of nodes leading to the given node 

25 starting from the root. 

26 """ 

27 ancestors = [] 

28 while node.parent is not None and isinstance(node.parent, Element): 

29 ancestors = [node.parent, *ancestors] 

30 node = node.parent 

31 

32 return ancestors 

33 

34 

35def path_names(node: Node) -> list[str]: 

36 """Get a list of nodes where each one is a child of 

37 the other leading to the node passed in. This gives a 

38 path to the node. 

39 

40 Does not include given node. 

41 

42 Args: 

43 node (Node): Node to find ancestors of. 

44 

45 Returns: 

46 list[str]: List of nodes leading to the given node 

47 starting from the root. 

48 """ 

49 ancestors = [] 

50 while node.parent is not None and isinstance(node.parent, Element): 

51 ancestors = [node.parent.tag, *ancestors] 

52 node = node.parent 

53 

54 return ancestors 

55 

56 

57def walk(node: Parent) -> Iterator: 

58 """Recursively traverse the node and it's chidlren as an iterator. 

59 Left to right depth first. 

60 """ 

61 

62 def get_children(n: Node) -> Iterator: 

63 yield n 

64 if isinstance(n, Parent): 

65 for child in n: 

66 yield from get_children(child) 

67 

68 yield node 

69 if isinstance(node, Parent): 

70 for child in node: 

71 yield from get_children(child) 

72 

73 

74def visit_all_after(start: Node) -> Iterator: 

75 """Recursively traverse the tree starting at given node.""" 

76 

77 def get_children(parent) -> Iterator: 

78 yield parent 

79 if parent.type in ["root", "element"]: 

80 for child in parent: 

81 yield from get_children(child) 

82 

83 parent = start.parent 

84 if parent is not None: 

85 idx = parent.index(start) 

86 if idx < len(parent) - 1: 

87 for child in parent[idx + 1 :]: 

88 yield from get_children(child)