Coverage for phml\utils\travel.py: 100%
33 statements
« prev ^ index » next coverage.py v6.5.0, created at 2022-12-08 11:07 -0600
« prev ^ index » next coverage.py v6.5.0, created at 2022-12-08 11:07 -0600
1"""utils.travel
3Collection of utilities that hep with traversing an ast or node tree.
4"""
6from typing import Iterator
8from phml.nodes import All_Nodes, Element, Root
10__all__ = ["path", "walk", "visit_children", "visit_all_after"]
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.
18 Does not include given node.
20 Args:
21 node (All_Nodes): Node to find ancestors of.
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
32 return ancestors
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 """
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)
46 yield node
47 if node.type in ["root", "element"]:
48 for child in visit_children(node):
49 yield from get_children(child)
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
58def visit_all_after(start: All_Nodes) -> Iterator:
59 """Recursively traverse the tree starting at given node."""
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)
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)