Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/lark/tree.py: 53%

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 from typing import Literal

14###{standalone

16class Meta:

18 empty: bool

19 line: int

20 column: int

21 start_pos: int

22 end_line: int

23 end_column: int

24 end_pos: int

25 orig_expansion: 'List[TerminalDef]'

26 match_tree: bool

28 def __init__(self):

29 self.empty = True

32_Leaf_T = TypeVar("_Leaf_T")

33Branch = Union[_Leaf_T, 'Tree[_Leaf_T]']

36class Tree(Generic[_Leaf_T]):

37 """The main tree class.

39 Creates a new tree, and stores "data" and "children" in attributes of the same name.

40 Trees can be hashed and compared.

42 Parameters:

43 data: The name of the rule or alias

44 children: List of matched sub-rules and terminals

45 meta: Line & Column numbers (if ``propagate_positions`` is enabled).

46 meta attributes: (line, column, end_line, end_column, start_pos, end_pos,

47 container_line, container_column, container_end_line, container_end_column)

48 container_* attributes consider all symbols, including those that have been inlined in the tree.

49 For example, in the rule 'a: _A B _C', the regular attributes will mark the start and end of B,

50 but the container_* attributes will also include _A and _C in the range. However, rules that

51 contain 'a' will consider it in full, including _A and _C for all attributes.

52 """

54 data: str

55 children: 'List[Branch[_Leaf_T]]'

57 def __init__(self, data: str, children: 'List[Branch[_Leaf_T]]', meta: Optional[Meta]=None) -> None:

58 self.data = data

59 self.children = children

60 self._meta = meta

62 @property

63 def meta(self) -> Meta:

64 if self._meta is None:

65 self._meta = Meta()

66 return self._meta

68 def __repr__(self):

69 return 'Tree(%r, %r)' % (self.data, self.children)

71 def _pretty_label(self):

72 return self.data

74 def _pretty(self, level, indent_str):

75 yield f'{indent_str*level}{self._pretty_label()}'

76 if len(self.children) == 1 and not isinstance(self.children[0], Tree):

77 yield f'\t{self.children[0]}\n'

78 else:

79 yield '\n'

80 for n in self.children:

81 if isinstance(n, Tree):

82 yield from n._pretty(level+1, indent_str)

83 else:

84 yield f'{indent_str*(level+1)}{n}\n'

86 def pretty(self, indent_str: str=' ') -> str:

87 """Returns an indented string representation of the tree.

89 Great for debugging.

90 """

91 return ''.join(self._pretty(0, indent_str))

93 def __rich__(self, parent:Optional['rich.tree.Tree']=None) -> 'rich.tree.Tree':

94 """Returns a tree widget for the 'rich' library.

96 Example:

97 ::

98 from rich import print

99 from lark import Tree

100

101 tree = Tree('root', ['node1', 'node2'])

102 print(tree)

103 """

104 return self._rich(parent)

105

106 def _rich(self, parent):

107 if parent:

108 tree = parent.add(f'[bold]{self.data}[/bold]')

109 else:

110 import rich.tree

111 tree = rich.tree.Tree(self.data)

112

113 for c in self.children:

114 if isinstance(c, Tree):

115 c._rich(tree)

116 else:

117 tree.add(f'[green]{c}[/green]')

118

119 return tree

120

121 def __eq__(self, other):

122 try:

123 return self.data == other.data and self.children == other.children

124 except AttributeError:

125 return False

126

127 def __ne__(self, other):

128 return not (self == other)

129

130 def __hash__(self) -> int:

131 return hash((self.data, tuple(self.children)))

132

133 def iter_subtrees(self) -> 'Iterator[Tree[_Leaf_T]]':

134 """Depth-first iteration.

135

136 Iterates over all the subtrees, never returning to the same node twice (Lark's parse-tree is actually a DAG).

137 """

138 queue = [self]

139 subtrees = dict()

140 for subtree in queue:

141 subtrees[id(subtree)] = subtree

142 queue += [c for c in reversed(subtree.children)

143 if isinstance(c, Tree) and id(c) not in subtrees]

144

145 del queue

146 return reversed(list(subtrees.values()))

147

148 def iter_subtrees_topdown(self):

149 """Breadth-first iteration.

150

151 Iterates over all the subtrees, return nodes in order like pretty() does.

152 """

153 stack = [self]

154 stack_append = stack.append

155 stack_pop = stack.pop

156 while stack:

157 node = stack_pop()

158 if not isinstance(node, Tree):

159 continue

160 yield node

161 for child in reversed(node.children):

162 stack_append(child)

163

164 def find_pred(self, pred: 'Callable[[Tree[_Leaf_T]], bool]') -> 'Iterator[Tree[_Leaf_T]]':

165 """Returns all nodes of the tree that evaluate pred(node) as true."""

166 return filter(pred, self.iter_subtrees())

167

168 def find_data(self, data: str) -> 'Iterator[Tree[_Leaf_T]]':

169 """Returns all nodes of the tree whose data equals the given data."""

170 return self.find_pred(lambda t: t.data == data)

171

172###}

173

174 def expand_kids_by_data(self, *data_values):

175 """Expand (inline) children with any of the given data values. Returns True if anything changed"""

176 changed = False

177 for i in range(len(self.children)-1, -1, -1):

178 child = self.children[i]

179 if isinstance(child, Tree) and child.data in data_values:

180 self.children[i:i+1] = child.children

181 changed = True

182 return changed

183

184

185 def scan_values(self, pred: 'Callable[[Branch[_Leaf_T]], bool]') -> Iterator[_Leaf_T]:

186 """Return all values in the tree that evaluate pred(value) as true.

187

188 This can be used to find all the tokens in the tree.

189

190 Example:

191 >>> all_tokens = tree.scan_values(lambda v: isinstance(v, Token))

192 """

193 for c in self.children:

194 if isinstance(c, Tree):

195 for t in c.scan_values(pred):

196 yield t

197 else:

198 if pred(c):

199 yield c

200

201 def __deepcopy__(self, memo):

202 return type(self)(self.data, deepcopy(self.children, memo), meta=self._meta)

203

204 def copy(self) -> 'Tree[_Leaf_T]':

205 return type(self)(self.data, self.children)

206

207 def set(self, data: str, children: 'List[Branch[_Leaf_T]]') -> None:

208 self.data = data

209 self.children = children

210

211

212ParseTree = Tree['Token']

213

214

215class SlottedTree(Tree):

216 __slots__ = 'data', 'children', 'rule', '_meta'

217

218

219def pydot__tree_to_png(tree: Tree, filename: str, rankdir: 'Literal["TB", "LR", "BT", "RL"]'="LR", **kwargs) -> None:

220 graph = pydot__tree_to_graph(tree, rankdir, **kwargs)

221 graph.write_png(filename)

222

223

224def pydot__tree_to_dot(tree: Tree, filename, rankdir="LR", **kwargs):

225 graph = pydot__tree_to_graph(tree, rankdir, **kwargs)

226 graph.write(filename)

227

228

229def pydot__tree_to_graph(tree: Tree, rankdir="LR", **kwargs):

230 """Creates a colorful image that represents the tree (data+children, without meta)

231

232 Possible values for `rankdir` are "TB", "LR", "BT", "RL", corresponding to

233 directed graphs drawn from top to bottom, from left to right, from bottom to

234 top, and from right to left, respectively.

235

236 `kwargs` can be any graph attribute (e. g. `dpi=200`). For a list of

237 possible attributes, see https://www.graphviz.org/doc/info/attrs.html.

238 """

239

240 import pydot # type: ignore[import-not-found]

241 graph = pydot.Dot(graph_type='digraph', rankdir=rankdir, **kwargs)

242

243 i = [0]

244

245 def new_leaf(leaf):

246 node = pydot.Node(i[0], label=repr(leaf))

247 i[0] += 1

248 graph.add_node(node)

249 return node

250

251 def _to_pydot(subtree):

252 color = hash(subtree.data) & 0xffffff

253 color |= 0x808080

254

255 subnodes = [_to_pydot(child) if isinstance(child, Tree) else new_leaf(child)

256 for child in subtree.children]

257 node = pydot.Node(i[0], style="filled", fillcolor="#%x" % color, label=subtree.data)

258 i[0] += 1

259 graph.add_node(node)

260

261 for subnode in subnodes:

262 graph.add_edge(pydot.Edge(node, subnode))

263

264 return node

265

266 _to_pydot(tree)

267 return graph