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

33 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-12-08 13:16 -0600

1"""utils.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 All_Nodes, Element, Root 

9 

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

11 

12 

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

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 (All_Nodes): Node to find ancestors of. 

22 

23 Returns: 

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

25 starting from the root. 

26 """ 

27 ancestors = [] 

28 while node.parent is not None: 

29 ancestors = [node.parent, *ancestors] 

30 node = node.parent 

31 

32 return ancestors 

33 

34 

35def walk(node: Root | Element) -> Iterator: 

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

37 Left to right depth first. 

38 """ 

39 

40 def get_children(parent) -> Iterator: 

41 yield parent 

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

43 for child in parent.children: 

44 yield from get_children(child) 

45 

46 yield node 

47 if node.type in ["root", "element"]: 

48 for child in visit_children(node): 

49 yield from get_children(child) 

50 

51 

52def visit_children(parent: Root | Element) -> Iterator: 

53 """Traverse the children of a Root or Element as an iterator.""" 

54 for child in parent.children: 

55 yield child 

56 

57 

58def visit_all_after(start: All_Nodes) -> Iterator: 

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

60 

61 def get_children(parent) -> Iterator: 

62 yield parent 

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

64 for child in parent.children: 

65 yield from get_children(child) 

66 

67 parent = start.parent 

68 idx = parent.children.index(start) 

69 if idx < len(parent.children) - 1: 

70 for child in parent.children[idx + 1 :]: 

71 yield from get_children(child)