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

1import sys 

2from copy import deepcopy 

3 

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

5 

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 

16 

17###{standalone 

18from collections import OrderedDict 

19 

20class Meta: 

21 

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 

31 

32 def __init__(self): 

33 self.empty = True 

34 

35 

36_Leaf_T = TypeVar("_Leaf_T") 

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

38 

39 

40class Tree(Generic[_Leaf_T]): 

41 """The main tree class. 

42 

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

44 Trees can be hashed and compared. 

45 

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 """ 

57 

58 data: str 

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

60 

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 

65 

66 @property 

67 def meta(self) -> Meta: 

68 if self._meta is None: 

69 self._meta = Meta() 

70 return self._meta 

71 

72 def __repr__(self): 

73 return 'Tree(%r, %r)' % (self.data, self.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 = 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] 

149 

150 del queue 

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

152 

153 def iter_subtrees_topdown(self): 

154 """Breadth-first iteration. 

155 

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) 

168 

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

172 

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) 

176 

177###} 

178 

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 

188 

189 

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. 

192 

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

194 

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 

205 

206 def __deepcopy__(self, memo): 

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

208 

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

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

211 

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

213 self.data = data 

214 self.children = children 

215 

216 

217ParseTree = Tree['Token'] 

218 

219 

220class SlottedTree(Tree): 

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

222 

223 

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) 

227 

228 

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

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

231 graph.write(filename) 

232 

233 

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

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

236 

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. 

240 

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 """ 

244 

245 import pydot # type: ignore[import] 

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

247 

248 i = [0] 

249 

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 

255 

256 def _to_pydot(subtree): 

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

258 color |= 0x808080 

259 

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) 

265 

266 for subnode in subnodes: 

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

268 

269 return node 

270 

271 _to_pydot(tree) 

272 return graph