try: from future_builtins import filter except ImportError: pass from copy import deepcopy ###{standalone from collections import OrderedDict class Meta: def __init__(self): self.empty = True class Tree(object): """The main tree class. Creates a new tree, and stores "data" and "children" in attributes of the same name. Trees can be hashed and compared. Parameters: data: The name of the rule or alias children: List of matched sub-rules and terminals meta: Line & Column numbers (if ``propagate_positions`` is enabled). meta attributes: line, column, start_pos, end_line, end_column, end_pos """ def __init__(self, data, children, meta=None): self.data = data self.children = children self._meta = meta @property def meta(self): if self._meta is None: self._meta = Meta() return self._meta def __repr__(self): return 'Tree(%r, %r)' % (self.data, self.children) def _pretty_label(self): return self.data def _pretty(self, level, indent_str): if len(self.children) == 1 and not isinstance(self.children[0], Tree): return [indent_str*level, self._pretty_label(), '\t', '%s' % (self.children[0],), '\n'] l = [indent_str*level, self._pretty_label(), '\n'] for n in self.children: if isinstance(n, Tree): l += n._pretty(level+1, indent_str) else: l += [indent_str*(level+1), '%s' % (n,), '\n'] return l def pretty(self, indent_str=' '): """Returns an indented string representation of the tree. Great for debugging. """ return ''.join(self._pretty(0, indent_str)) def __eq__(self, other): try: return self.data == other.data and self.children == other.children except AttributeError: return False def __ne__(self, other): return not (self == other) def __hash__(self): return hash((self.data, tuple(self.children))) def iter_subtrees(self): """Depth-first iteration. Iterates over all the subtrees, never returning to the same node twice (Lark's parse-tree is actually a DAG). """ queue = [self] subtrees = OrderedDict() for subtree in queue: subtrees[id(subtree)] = subtree queue += [c for c in reversed(subtree.children) if isinstance(c, Tree) and id(c) not in subtrees] del queue return reversed(list(subtrees.values())) def find_pred(self, pred): """Returns all nodes of the tree that evaluate pred(node) as true.""" return filter(pred, self.iter_subtrees()) def find_data(self, data): """Returns all nodes of the tree whose data equals the given data.""" return self.find_pred(lambda t: t.data == data) ###} def expand_kids_by_index(self, *indices): """Expand (inline) children at the given indices""" for i in sorted(indices, reverse=True): # reverse so that changing tail won't affect indices kid = self.children[i] self.children[i:i+1] = kid.children def scan_values(self, pred): for c in self.children: if isinstance(c, Tree): for t in c.scan_values(pred): yield t else: if pred(c): yield c def iter_subtrees_topdown(self): """Breadth-first iteration. Iterates over all the subtrees, return nodes in order like pretty() does. """ stack = [self] while stack: node = stack.pop() if not isinstance(node, Tree): continue yield node for n in reversed(node.children): stack.append(n) def __deepcopy__(self, memo): return type(self)(self.data, deepcopy(self.children, memo), meta=self._meta) def copy(self): return type(self)(self.data, self.children) def set(self, data, children): self.data = data self.children = children # XXX Deprecated! Here for backwards compatibility <0.6.0 @property def line(self): return self.meta.line @property def column(self): return self.meta.column @property def end_line(self): return self.meta.end_line @property def end_column(self): return self.meta.end_column class SlottedTree(Tree): __slots__ = 'data', 'children', 'rule', '_meta' def pydot__tree_to_png(tree, filename, rankdir="LR", **kwargs): graph = pydot__tree_to_graph(tree, rankdir, **kwargs) graph.write_png(filename) def pydot__tree_to_dot(tree, filename, rankdir="LR", **kwargs): graph = pydot__tree_to_graph(tree, rankdir, **kwargs) graph.write(filename) def pydot__tree_to_graph(tree, rankdir="LR", **kwargs): """Creates a colorful image that represents the tree (data+children, without meta) Possible values for `rankdir` are "TB", "LR", "BT", "RL", corresponding to directed graphs drawn from top to bottom, from left to right, from bottom to top, and from right to left, respectively. `kwargs` can be any graph attribute (e. g. `dpi=200`). For a list of possible attributes, see https://www.graphviz.org/doc/info/attrs.html. """ import pydot graph = pydot.Dot(graph_type='digraph', rankdir=rankdir, **kwargs) i = [0] def new_leaf(leaf): node = pydot.Node(i[0], label=repr(leaf)) i[0] += 1 graph.add_node(node) return node def _to_pydot(subtree): color = hash(subtree.data) & 0xffffff color |= 0x808080 subnodes = [_to_pydot(child) if isinstance(child, Tree) else new_leaf(child) for child in subtree.children] node = pydot.Node(i[0], style="filled", fillcolor="#%x" % color, label=subtree.data) i[0] += 1 graph.add_node(node) for subnode in subnodes: graph.add_edge(pydot.Edge(node, subnode)) return node _to_pydot(tree) return graph