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

1import sys 

2from copy import deepcopy 

3 

4from typing import List, Callable, Iterator, Union, Optional, Generic, TypeVar, TYPE_CHECKING 

5 

6from .lexer import Token 

7 

8if TYPE_CHECKING: 

9 from .lexer import TerminalDef 

10 try: 

11 import rich 

12 except ImportError: 

13 pass 

14 from typing import Literal 

15 

16###{standalone 

17 

18class Meta: 

19 

20 empty: bool 

21 line: int 

22 column: int 

23 start_pos: int 

24 end_line: int 

25 end_column: int 

26 end_pos: int 

27 orig_expansion: 'List[TerminalDef]' 

28 match_tree: bool 

29 

30 def __init__(self): 

31 self.empty = True 

32 

33 

34_Leaf_T = TypeVar("_Leaf_T") 

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

36 

37 

38class Tree(Generic[_Leaf_T]): 

39 """The main tree class. 

40 

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

42 Trees can be hashed and compared. 

43 

44 Parameters: 

45 data: The name of the rule or alias 

46 children: List of matched sub-rules and terminals 

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

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

49 container_line, container_column, container_end_line, container_end_column) 

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

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

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

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

54 """ 

55 

56 data: str 

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

58 

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

60 self.data = data 

61 self.children = children 

62 self._meta = meta 

63 

64 @property 

65 def meta(self) -> Meta: 

66 if self._meta is None: 

67 self._meta = Meta() 

68 return self._meta 

69 

70 def __repr__(self): 

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

72 

73 __match_args__ = ("data", "children") 

74 

75 def _pretty_label(self): 

76 return self.data 

77 

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' 

89 

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

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

92 

93 Great for debugging. 

94 """ 

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

96 

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

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

99 

100 Example: 

101 :: 

102 from rich import print 

103 from lark import Tree 

104 

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

106 print(tree) 

107 """ 

108 return self._rich(parent) 

109 

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) 

116 

117 for c in self.children: 

118 if isinstance(c, Tree): 

119 c._rich(tree) 

120 else: 

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

122 

123 return tree 

124 

125 def __eq__(self, other): 

126 try: 

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

128 except AttributeError: 

129 return False 

130 

131 def __ne__(self, other): 

132 return not (self == other) 

133 

134 def __hash__(self) -> int: 

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

136 

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

138 """Depth-first iteration. 

139 

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 = dict() 

144 for subtree in queue: 

145 subtrees[id(subtree)] = subtree 

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

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

148 

149 del queue 

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

151 

152 def iter_subtrees_topdown(self): 

153 """Breadth-first iteration. 

154 

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

156 """ 

157 stack = [self] 

158 stack_append = stack.append 

159 stack_pop = stack.pop 

160 while stack: 

161 node = stack_pop() 

162 if not isinstance(node, Tree): 

163 continue 

164 yield node 

165 for child in reversed(node.children): 

166 stack_append(child) 

167 

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

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

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

171 

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

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

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

175 

176###} 

177 

178 def find_token(self, token_type: str) -> Iterator[_Leaf_T]: 

179 """Returns all tokens whose type equals the given token_type. 

180 

181 This is a recursive function that will find tokens in all the subtrees. 

182 

183 Example: 

184 >>> term_tokens = tree.find_token('TERM') 

185 """ 

186 return self.scan_values(lambda v: isinstance(v, Token) and v.type == token_type) 

187 

188 def expand_kids_by_data(self, *data_values): 

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

190 changed = False 

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

192 child = self.children[i] 

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

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

195 changed = True 

196 return changed 

197 

198 

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

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

201 

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

203 

204 Example: 

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

206 """ 

207 for c in self.children: 

208 if isinstance(c, Tree): 

209 for t in c.scan_values(pred): 

210 yield t 

211 else: 

212 if pred(c): 

213 yield c 

214 

215 def __deepcopy__(self, memo): 

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

217 

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

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

220 

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

222 self.data = data 

223 self.children = children 

224 

225 

226ParseTree = Tree['Token'] 

227 

228 

229class SlottedTree(Tree): 

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

231 

232 

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

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

235 graph.write_png(filename) 

236 

237 

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

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

240 graph.write(filename) 

241 

242 

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

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

245 

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

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

248 top, and from right to left, respectively. 

249 

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

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

252 """ 

253 

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

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

256 

257 i = [0] 

258 

259 def new_leaf(leaf): 

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

261 i[0] += 1 

262 graph.add_node(node) 

263 return node 

264 

265 def _to_pydot(subtree): 

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

267 color |= 0x808080 

268 

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

270 for child in subtree.children] 

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

272 i[0] += 1 

273 graph.add_node(node) 

274 

275 for subnode in subnodes: 

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

277 

278 return node 

279 

280 _to_pydot(tree) 

281 return graph