codekingpro/portable-devtools

1from typing import Iterator, List, Optional, Tuple2 3from ._loop import loop_first, loop_last4from .console import Console, ConsoleOptions, RenderableType, RenderResult5from .jupyter import JupyterMixin6from .measure import Measurement7from .segment import Segment8from .style import Style, StyleStack, StyleType9from .styled import Styled10 11 12class Tree(JupyterMixin):13    """A renderable for a tree structure.14 15    Args:16        label (RenderableType): The renderable or str for the tree label.17        style (StyleType, optional): Style of this tree. Defaults to "tree".18        guide_style (StyleType, optional): Style of the guide lines. Defaults to "tree.line".19        expanded (bool, optional): Also display children. Defaults to True.20        highlight (bool, optional): Highlight renderable (if str). Defaults to False.21    """22 23    def __init__(24        self,25        label: RenderableType,26        *,27        style: StyleType = "tree",28        guide_style: StyleType = "tree.line",29        expanded: bool = True,30        highlight: bool = False,31        hide_root: bool = False,32    ) -> None:33        self.label = label34        self.style = style35        self.guide_style = guide_style36        self.children: List[Tree] = []37        self.expanded = expanded38        self.highlight = highlight39        self.hide_root = hide_root40 41    def add(42        self,43        label: RenderableType,44        *,45        style: Optional[StyleType] = None,46        guide_style: Optional[StyleType] = None,47        expanded: bool = True,48        highlight: Optional[bool] = False,49    ) -> "Tree":50        """Add a child tree.51 52        Args:53            label (RenderableType): The renderable or str for the tree label.54            style (StyleType, optional): Style of this tree. Defaults to "tree".55            guide_style (StyleType, optional): Style of the guide lines. Defaults to "tree.line".56            expanded (bool, optional): Also display children. Defaults to True.57            highlight (Optional[bool], optional): Highlight renderable (if str). Defaults to False.58 59        Returns:60            Tree: A new child Tree, which may be further modified.61        """62        node = Tree(63            label,64            style=self.style if style is None else style,65            guide_style=self.guide_style if guide_style is None else guide_style,66            expanded=expanded,67            highlight=self.highlight if highlight is None else highlight,68        )69        self.children.append(node)70        return node71 72    def __rich_console__(73        self, console: "Console", options: "ConsoleOptions"74    ) -> "RenderResult":75        stack: List[Iterator[Tuple[bool, Tree]]] = []76        pop = stack.pop77        push = stack.append78        new_line = Segment.line()79 80        get_style = console.get_style81        null_style = Style.null()82        guide_style = get_style(self.guide_style, default="") or null_style83        SPACE, CONTINUE, FORK, END = range(4)84 85        ASCII_GUIDES = ("    ", "|   ", "+-- ", "`-- ")86        TREE_GUIDES = [87            ("    ", "│   ", "├── ", "└── "),88            ("    ", "┃   ", "┣━━ ", "┗━━ "),89            ("    ", "║   ", "╠══ ", "╚══ "),90        ]91        _Segment = Segment92 93        def make_guide(index: int, style: Style) -> Segment:94            """Make a Segment for a level of the guide lines."""95            if options.ascii_only:96                line = ASCII_GUIDES[index]97            else:98                guide = 1 if style.bold else (2 if style.underline2 else 0)99                line = TREE_GUIDES[0 if options.legacy_windows else guide][index]100            return _Segment(line, style)101 102        levels: List[Segment] = [make_guide(CONTINUE, guide_style)]103        push(iter(loop_last([self])))104 105        guide_style_stack = StyleStack(get_style(self.guide_style))106        style_stack = StyleStack(get_style(self.style))107        remove_guide_styles = Style(bold=False, underline2=False)108 109        depth = 0110 111        while stack:112            stack_node = pop()113            try:114                last, node = next(stack_node)115            except StopIteration:116                levels.pop()117                if levels:118                    guide_style = levels[-1].style or null_style119                    levels[-1] = make_guide(FORK, guide_style)120                    guide_style_stack.pop()121                    style_stack.pop()122                continue123            push(stack_node)124            if last:125                levels[-1] = make_guide(END, levels[-1].style or null_style)126 127            guide_style = guide_style_stack.current + get_style(node.guide_style)128            style = style_stack.current + get_style(node.style)129            prefix = levels[(2 if self.hide_root else 1) :]130            renderable_lines = console.render_lines(131                Styled(node.label, style),132                options.update(133                    width=options.max_width134                    - sum(level.cell_length for level in prefix),135                    highlight=self.highlight,136                    height=None,137                ),138                pad=options.justify is not None,139            )140 141            if not (depth == 0 and self.hide_root):142                for first, line in loop_first(renderable_lines):143                    if prefix:144                        yield from _Segment.apply_style(145                            prefix,146                            style.background_style,147                            post_style=remove_guide_styles,148                        )149                    yield from line150                    yield new_line151                    if first and prefix:152                        prefix[-1] = make_guide(153                            SPACE if last else CONTINUE, prefix[-1].style or null_style154                        )155 156            if node.expanded and node.children:157                levels[-1] = make_guide(158                    SPACE if last else CONTINUE, levels[-1].style or null_style159                )160                levels.append(161                    make_guide(END if len(node.children) == 1 else FORK, guide_style)162                )163                style_stack.push(get_style(node.style))164                guide_style_stack.push(get_style(node.guide_style))165                push(iter(loop_last(node.children)))166                depth += 1167 168    def __rich_measure__(169        self, console: "Console", options: "ConsoleOptions"170    ) -> "Measurement":171        stack: List[Iterator[Tree]] = [iter([self])]172        pop = stack.pop173        push = stack.append174        minimum = 0175        maximum = 0176        measure = Measurement.get177        level = 0178        while stack:179            iter_tree = pop()180            try:181                tree = next(iter_tree)182            except StopIteration:183                level -= 1184                continue185            push(iter_tree)186            min_measure, max_measure = measure(console, options, tree.label)187            indent = level * 4188            minimum = max(min_measure + indent, minimum)189            maximum = max(max_measure + indent, maximum)190            if tree.expanded and tree.children:191                push(iter(tree.children))192                level += 1193        return Measurement(minimum, maximum)194 195 196if __name__ == "__main__":  # pragma: no cover197    from rich.console import Group198    from rich.markdown import Markdown199    from rich.panel import Panel200    from rich.syntax import Syntax201    from rich.table import Table202 203    table = Table(row_styles=["", "dim"])204 205    table.add_column("Released", style="cyan", no_wrap=True)206    table.add_column("Title", style="magenta")207    table.add_column("Box Office", justify="right", style="green")208 209    table.add_row("Dec 20, 2019", "Star Wars: The Rise of Skywalker", "$952,110,690")210    table.add_row("May 25, 2018", "Solo: A Star Wars Story", "$393,151,347")211    table.add_row("Dec 15, 2017", "Star Wars Ep. V111: The Last Jedi", "$1,332,539,889")212    table.add_row("Dec 16, 2016", "Rogue One: A Star Wars Story", "$1,332,439,889")213 214    code = """\215class Segment(NamedTuple):216    text: str = ""217    style: Optional[Style] = None218    is_control: bool = False219"""220    syntax = Syntax(code, "python", theme="monokai", line_numbers=True)221 222    markdown = Markdown(223        """\224### example.md225> Hello, World!226>227> Markdown _all_ the things228"""229    )230 231    root = Tree("🌲 [b green]Rich Tree", highlight=True, hide_root=True)232 233    node = root.add(":file_folder: Renderables", guide_style="red")234    simple_node = node.add(":file_folder: [bold yellow]Atomic", guide_style="uu green")235    simple_node.add(Group("📄 Syntax", syntax))236    simple_node.add(Group("📄 Markdown", Panel(markdown, border_style="green")))237 238    containers_node = node.add(239        ":file_folder: [bold magenta]Containers", guide_style="bold magenta"240    )241    containers_node.expanded = True242    panel = Panel.fit("Just a panel", border_style="red")243    containers_node.add(Group("📄 Panels", panel))244 245    containers_node.add(Group("📄 [b magenta]Table", table))246 247    console = Console()248 249    console.print(root)250