diff --git a/docs/classes.rst b/docs/classes.rst index 2a1c7706..6e88fae4 100644 --- a/docs/classes.rst +++ b/docs/classes.rst @@ -39,7 +39,7 @@ Tree .. autoclass:: lark.Tree :members: pretty, find_pred, find_data, iter_subtrees, scan_values, - iter_subtrees_topdown + iter_subtrees_topdown, __rich__ Token ----- diff --git a/lark/tree.py b/lark/tree.py index cea5aaa9..51c962d1 100644 --- a/lark/tree.py +++ b/lark/tree.py @@ -86,6 +86,34 @@ def pretty(self, indent_str: str=' ') -> str: """ return ''.join(self._pretty(0, indent_str)) + def __rich__(self, parent:'rich.tree.Tree'=None) -> 'rich.tree.Tree': + """Returns a tree widget for the 'rich' library. + + Example: + :: + from rich import print + from lark import Tree + + tree = Tree('root', ['node1', 'node2']) + print(tree) + """ + return self._rich(parent) + + def _rich(self, parent): + if parent: + tree = parent.add(f'[bold]{self.data}[/bold]') + else: + import rich.tree + tree = rich.tree.Tree(self.data) + + for c in self.children: + if isinstance(c, Tree): + c._rich(tree) + else: + tree.add(f'[green]{c}[/green]') + + return tree + def __eq__(self, other): try: return self.data == other.data and self.children == other.children