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
« prev ^ index » next coverage.py v6.5.0, created at 2023-04-12 14:26 -0500
1"""utilities.travel
3Collection of utilities that hep with traversing an ast or node tree.
4"""
6from typing import Iterator
8from phml.nodes import Element, Node, Parent
10__all__ = ["path", "path_names", "walk", "visit_all_after"]
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.
18 Does not include given node.
20 Args:
21 node (Node): Node to find ancestors of.
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
32 return ancestors
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.
40 Does not include given node.
42 Args:
43 node (Node): Node to find ancestors of.
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
54 return ancestors
57def walk(node: Parent) -> Iterator:
58 """Recursively traverse the node and it's chidlren as an iterator.
59 Left to right depth first.
60 """
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)
68 yield node
69 if isinstance(node, Parent):
70 for child in node:
71 yield from get_children(child)
74def visit_all_after(start: Node) -> Iterator:
75 """Recursively traverse the tree starting at given node."""
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)
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)