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 11GuideType = Tuple[str, str, str, str]12 13 14class Tree(JupyterMixin):15    """A renderable for a tree structure.16 17    Attributes:18        ASCII_GUIDES (GuideType): Guide lines used when Console.ascii_only is True.19        TREE_GUIDES (List[GuideType, GuideType, GuideType]): Default guide lines.20 21    Args:22        label (RenderableType): The renderable or str for the tree label.23        style (StyleType, optional): Style of this tree. Defaults to "tree".24        guide_style (StyleType, optional): Style of the guide lines. Defaults to "tree.line".25        expanded (bool, optional): Also display children. Defaults to True.26        highlight (bool, optional): Highlight renderable (if str). Defaults to False.27        hide_root (bool, optional): Hide the root node. Defaults to False.28    """29 30    ASCII_GUIDES = ("    ", "|   ", "+-- ", "`-- ")31    TREE_GUIDES = [32        ("    ", "│   ", "├── ", "└── "),33        ("    ", "┃   ", "┣━━ ", "┗━━ "),34        ("    ", "║   ", "╠══ ", "╚══ "),35    ]36 37    def __init__(38        self,39        label: RenderableType,40        *,41        style: StyleType = "tree",42        guide_style: StyleType = "tree.line",43        expanded: bool = True,44        highlight: bool = False,45        hide_root: bool = False,46    ) -> None:47        self.label = label48        self.style = style49        self.guide_style = guide_style50        self.children: List[Tree] = []51        self.expanded = expanded52        self.highlight = highlight53        self.hide_root = hide_root54 55    def add(56        self,57        label: RenderableType,58        *,59        style: Optional[StyleType] = None,60        guide_style: Optional[StyleType] = None,61        expanded: bool = True,62        highlight: Optional[bool] = False,63    ) -> "Tree":64        """Add a child tree.65 66        Args:67            label (RenderableType): The renderable or str for the tree label.68            style (StyleType, optional): Style of this tree. Defaults to "tree".69            guide_style (StyleType, optional): Style of the guide lines. Defaults to "tree.line".70            expanded (bool, optional): Also display children. Defaults to True.71            highlight (Optional[bool], optional): Highlight renderable (if str). Defaults to False.72 73        Returns:74            Tree: A new child Tree, which may be further modified.75        """76        node = Tree(77            label,78            style=self.style if style is None else style,79            guide_style=self.guide_style if guide_style is None else guide_style,80            expanded=expanded,81            highlight=self.highlight if highlight is None else highlight,82        )83        self.children.append(node)84        return node85 86    def __rich_console__(87        self, console: "Console", options: "ConsoleOptions"88    ) -> "RenderResult":89        stack: List[Iterator[Tuple[bool, Tree]]] = []90        pop = stack.pop91        push = stack.append92        new_line = Segment.line()93 94        get_style = console.get_style95        null_style = Style.null()96        guide_style = get_style(self.guide_style, default="") or null_style97        SPACE, CONTINUE, FORK, END = range(4)98 99        _Segment = Segment100 101        def make_guide(index: int, style: Style) -> Segment:102            """Make a Segment for a level of the guide lines."""103            if options.ascii_only:104                line = self.ASCII_GUIDES[index]105            else:106                guide = 1 if style.bold else (2 if style.underline2 else 0)107                line = self.TREE_GUIDES[0 if options.legacy_windows else guide][index]108            return _Segment(line, style)109 110        levels: List[Segment] = [make_guide(CONTINUE, guide_style)]111        push(iter(loop_last([self])))112 113        guide_style_stack = StyleStack(get_style(self.guide_style))114        style_stack = StyleStack(get_style(self.style))115        remove_guide_styles = Style(bold=False, underline2=False)116 117        depth = 0118 119        while stack:120            stack_node = pop()121            try:122                last, node = next(stack_node)123            except StopIteration:124                levels.pop()125                if levels:126                    guide_style = levels[-1].style or null_style127                    levels[-1] = make_guide(FORK, guide_style)128                    guide_style_stack.pop()129                    style_stack.pop()130                continue131            push(stack_node)132            if last:133                levels[-1] = make_guide(END, levels[-1].style or null_style)134 135            guide_style = guide_style_stack.current + get_style(node.guide_style)136            style = style_stack.current + get_style(node.style)137            prefix = levels[(2 if self.hide_root else 1) :]138            renderable_lines = console.render_lines(139                Styled(node.label, style),140                options.update(141                    width=options.max_width142                    - sum(level.cell_length for level in prefix),143                    highlight=self.highlight,144                    height=None,145                ),146                pad=options.justify is not None,147            )148 149            if not (depth == 0 and self.hide_root):150                for first, line in loop_first(renderable_lines):151                    if prefix:152                        yield from _Segment.apply_style(153                            prefix,154                            style.background_style,155                            post_style=remove_guide_styles,156                        )157                    yield from line158                    yield new_line159                    if first and prefix:160                        prefix[-1] = make_guide(161                            SPACE if last else CONTINUE, prefix[-1].style or null_style162                        )163 164            if node.expanded and node.children:165                levels[-1] = make_guide(166                    SPACE if last else CONTINUE, levels[-1].style or null_style167                )168                levels.append(169                    make_guide(END if len(node.children) == 1 else FORK, guide_style)170                )171                style_stack.push(get_style(node.style))172                guide_style_stack.push(get_style(node.guide_style))173                push(iter(loop_last(node.children)))174                depth += 1175 176    def __rich_measure__(177        self, console: "Console", options: "ConsoleOptions"178    ) -> "Measurement":179        stack: List[Iterator[Tree]] = [iter([self])]180        pop = stack.pop181        push = stack.append182        minimum = 0183        maximum = 0184        measure = Measurement.get185        level = 0186        while stack:187            iter_tree = pop()188            try:189                tree = next(iter_tree)190            except StopIteration:191                level -= 1192                continue193            push(iter_tree)194            min_measure, max_measure = measure(console, options, tree.label)195            indent = level * 4196            minimum = max(min_measure + indent, minimum)197            maximum = max(max_measure + indent, maximum)198            if tree.expanded and tree.children:199                push(iter(tree.children))200                level += 1201        return Measurement(minimum, maximum)202 203 204if __name__ == "__main__":  # pragma: no cover205    from rich.console import Group206    from rich.markdown import Markdown207    from rich.panel import Panel208    from rich.syntax import Syntax209    from rich.table import Table210 211    table = Table(row_styles=["", "dim"])212 213    table.add_column("Released", style="cyan", no_wrap=True)214    table.add_column("Title", style="magenta")215    table.add_column("Box Office", justify="right", style="green")216 217    table.add_row("Dec 20, 2019", "Star Wars: The Rise of Skywalker", "$952,110,690")218    table.add_row("May 25, 2018", "Solo: A Star Wars Story", "$393,151,347")219    table.add_row("Dec 15, 2017", "Star Wars Ep. V111: The Last Jedi", "$1,332,539,889")220    table.add_row("Dec 16, 2016", "Rogue One: A Star Wars Story", "$1,332,439,889")221 222    code = """\223class Segment(NamedTuple):224    text: str = ""225    style: Optional[Style] = None226    is_control: bool = False227"""228    syntax = Syntax(code, "python", theme="monokai", line_numbers=True)229 230    markdown = Markdown(231        """\232### example.md233> Hello, World!234>235> Markdown _all_ the things236"""237    )238 239    root = Tree("🌲 [b green]Rich Tree", highlight=True, hide_root=True)240 241    node = root.add(":file_folder: Renderables", guide_style="red")242    simple_node = node.add(":file_folder: [bold yellow]Atomic", guide_style="uu green")243    simple_node.add(Group("📄 Syntax", syntax))244    simple_node.add(Group("📄 Markdown", Panel(markdown, border_style="green")))245 246    containers_node = node.add(247        ":file_folder: [bold magenta]Containers", guide_style="bold magenta"248    )249    containers_node.expanded = True250    panel = Panel.fit("Just a panel", border_style="red")251    containers_node.add(Group("📄 Panels", panel))252 253    containers_node.add(Group("📄 [b magenta]Table", table))254 255    console = Console()256 257    console.print(root)258