Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/lark/tree.py: 55%
150 statements
« prev ^ index » next coverage.py v7.3.1, created at 2023-09-25 06:30 +0000
« prev ^ index » next coverage.py v7.3.1, created at 2023-09-25 06:30 +0000
1import sys
2from copy import deepcopy
4from typing import List, Callable, Iterator, Union, Optional, Generic, TypeVar, TYPE_CHECKING
6if TYPE_CHECKING:
7 from .lexer import TerminalDef, Token
8 try:
9 import rich
10 except ImportError:
11 pass
12 if sys.version_info >= (3, 8):
13 from typing import Literal
14 else:
15 from typing_extensions import Literal
17###{standalone
18from collections import OrderedDict
20class Meta:
22 empty: bool
23 line: int
24 column: int
25 start_pos: int
26 end_line: int
27 end_column: int
28 end_pos: int
29 orig_expansion: 'List[TerminalDef]'
30 match_tree: bool
32 def __init__(self):
33 self.empty = True
36_Leaf_T = TypeVar("_Leaf_T")
37Branch = Union[_Leaf_T, 'Tree[_Leaf_T]']
40class Tree(Generic[_Leaf_T]):
41 """The main tree class.
43 Creates a new tree, and stores "data" and "children" in attributes of the same name.
44 Trees can be hashed and compared.
46 Parameters:
47 data: The name of the rule or alias
48 children: List of matched sub-rules and terminals
49 meta: Line & Column numbers (if ``propagate_positions`` is enabled).
50 meta attributes: (line, column, end_line, end_column, start_pos, end_pos,
51 container_line, container_column, container_end_line, container_end_column)
52 container_* attributes consider all symbols, including those that have been inlined in the tree.
53 For example, in the rule 'a: _A B _C', the regular attributes will mark the start and end of B,
54 but the container_* attributes will also include _A and _C in the range. However, rules that
55 contain 'a' will consider it in full, including _A and _C for all attributes.
56 """
58 data: str
59 children: 'List[Branch[_Leaf_T]]'
61 def __init__(self, data: str, children: 'List[Branch[_Leaf_T]]', meta: Optional[Meta]=None) -> None:
62 self.data = data
63 self.children = children
64 self._meta = meta
66 @property
67 def meta(self) -> Meta:
68 if self._meta is None:
69 self._meta = Meta()
70 return self._meta
72 def __repr__(self):
73 return 'Tree(%r, %r)' % (self.data, self.children)
75 def _pretty_label(self):
76 return self.data
78 def _pretty(self, level, indent_str):
79 yield f'{indent_str*level}{self._pretty_label()}'
80 if len(self.children) == 1 and not isinstance(self.children[0], Tree):
81 yield f'\t{self.children[0]}\n'
82 else:
83 yield '\n'
84 for n in self.children:
85 if isinstance(n, Tree):
86 yield from n._pretty(level+1, indent_str)
87 else:
88 yield f'{indent_str*(level+1)}{n}\n'
90 def pretty(self, indent_str: str=' ') -> str:
91 """Returns an indented string representation of the tree.
93 Great for debugging.
94 """
95 return ''.join(self._pretty(0, indent_str))
97 def __rich__(self, parent:Optional['rich.tree.Tree']=None) -> 'rich.tree.Tree':
98 """Returns a tree widget for the 'rich' library.
100 Example:
101 ::
102 from rich import print
103 from lark import Tree
105 tree = Tree('root', ['node1', 'node2'])
106 print(tree)
107 """
108 return self._rich(parent)
110 def _rich(self, parent):
111 if parent:
112 tree = parent.add(f'[bold]{self.data}[/bold]')
113 else:
114 import rich.tree
115 tree = rich.tree.Tree(self.data)
117 for c in self.children:
118 if isinstance(c, Tree):
119 c._rich(tree)
120 else:
121 tree.add(f'[green]{c}[/green]')
123 return tree
125 def __eq__(self, other):
126 try:
127 return self.data == other.data and self.children == other.children
128 except AttributeError:
129 return False
131 def __ne__(self, other):
132 return not (self == other)
134 def __hash__(self) -> int:
135 return hash((self.data, tuple(self.children)))
137 def iter_subtrees(self) -> 'Iterator[Tree[_Leaf_T]]':
138 """Depth-first iteration.
140 Iterates over all the subtrees, never returning to the same node twice (Lark's parse-tree is actually a DAG).
141 """
142 queue = [self]
143 subtrees = OrderedDict()
144 for subtree in queue:
145 subtrees[id(subtree)] = subtree
146 # Reason for type ignore https://github.com/python/mypy/issues/10999
147 queue += [c for c in reversed(subtree.children) # type: ignore[misc]
148 if isinstance(c, Tree) and id(c) not in subtrees]
150 del queue
151 return reversed(list(subtrees.values()))
153 def iter_subtrees_topdown(self):
154 """Breadth-first iteration.
156 Iterates over all the subtrees, return nodes in order like pretty() does.
157 """
158 stack = [self]
159 stack_append = stack.append
160 stack_pop = stack.pop
161 while stack:
162 node = stack_pop()
163 if not isinstance(node, Tree):
164 continue
165 yield node
166 for child in reversed(node.children):
167 stack_append(child)
169 def find_pred(self, pred: 'Callable[[Tree[_Leaf_T]], bool]') -> 'Iterator[Tree[_Leaf_T]]':
170 """Returns all nodes of the tree that evaluate pred(node) as true."""
171 return filter(pred, self.iter_subtrees())
173 def find_data(self, data: str) -> 'Iterator[Tree[_Leaf_T]]':
174 """Returns all nodes of the tree whose data equals the given data."""
175 return self.find_pred(lambda t: t.data == data)
177###}
179 def expand_kids_by_data(self, *data_values):
180 """Expand (inline) children with any of the given data values. Returns True if anything changed"""
181 changed = False
182 for i in range(len(self.children)-1, -1, -1):
183 child = self.children[i]
184 if isinstance(child, Tree) and child.data in data_values:
185 self.children[i:i+1] = child.children
186 changed = True
187 return changed
190 def scan_values(self, pred: 'Callable[[Branch[_Leaf_T]], bool]') -> Iterator[_Leaf_T]:
191 """Return all values in the tree that evaluate pred(value) as true.
193 This can be used to find all the tokens in the tree.
195 Example:
196 >>> all_tokens = tree.scan_values(lambda v: isinstance(v, Token))
197 """
198 for c in self.children:
199 if isinstance(c, Tree):
200 for t in c.scan_values(pred):
201 yield t
202 else:
203 if pred(c):
204 yield c
206 def __deepcopy__(self, memo):
207 return type(self)(self.data, deepcopy(self.children, memo), meta=self._meta)
209 def copy(self) -> 'Tree[_Leaf_T]':
210 return type(self)(self.data, self.children)
212 def set(self, data: str, children: 'List[Branch[_Leaf_T]]') -> None:
213 self.data = data
214 self.children = children
217ParseTree = Tree['Token']
220class SlottedTree(Tree):
221 __slots__ = 'data', 'children', 'rule', '_meta'
224def pydot__tree_to_png(tree: Tree, filename: str, rankdir: 'Literal["TB", "LR", "BT", "RL"]'="LR", **kwargs) -> None:
225 graph = pydot__tree_to_graph(tree, rankdir, **kwargs)
226 graph.write_png(filename)
229def pydot__tree_to_dot(tree: Tree, filename, rankdir="LR", **kwargs):
230 graph = pydot__tree_to_graph(tree, rankdir, **kwargs)
231 graph.write(filename)
234def pydot__tree_to_graph(tree: Tree, rankdir="LR", **kwargs):
235 """Creates a colorful image that represents the tree (data+children, without meta)
237 Possible values for `rankdir` are "TB", "LR", "BT", "RL", corresponding to
238 directed graphs drawn from top to bottom, from left to right, from bottom to
239 top, and from right to left, respectively.
241 `kwargs` can be any graph attribute (e. g. `dpi=200`). For a list of
242 possible attributes, see https://www.graphviz.org/doc/info/attrs.html.
243 """
245 import pydot # type: ignore[import]
246 graph = pydot.Dot(graph_type='digraph', rankdir=rankdir, **kwargs)
248 i = [0]
250 def new_leaf(leaf):
251 node = pydot.Node(i[0], label=repr(leaf))
252 i[0] += 1
253 graph.add_node(node)
254 return node
256 def _to_pydot(subtree):
257 color = hash(subtree.data) & 0xffffff
258 color |= 0x808080
260 subnodes = [_to_pydot(child) if isinstance(child, Tree) else new_leaf(child)
261 for child in subtree.children]
262 node = pydot.Node(i[0], style="filled", fillcolor="#%x" % color, label=subtree.data)
263 i[0] += 1
264 graph.add_node(node)
266 for subnode in subnodes:
267 graph.add_edge(pydot.Edge(node, subnode))
269 return node
271 _to_pydot(tree)
272 return graph