venv/lib/python3.11/site-packages/pip/_vendor/rich/tree.py

   1 from typing import Iterator, List, Optional, Tuple
   2
   3 from ._loop import loop_first, loop_last
   4 from .console import Console, ConsoleOptions, RenderableType, RenderResult
   5 from .jupyter import JupyterMixin
   6 from .measure import Measurement
   7 from .segment import Segment
   8 from .style import Style, StyleStack, StyleType
   9 from .styled import Styled
  10
  11
  12 class 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 = label
  34         self.style = style
  35         self.guide_style = guide_style
  36         self.children: List[Tree] = []
  37         self.expanded = expanded
  38         self.highlight = highlight
  39         self.hide_root = hide_root
  40
  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 node
  71
  72     def __rich_console__(
  73         self, console: "Console", options: "ConsoleOptions"
  74     ) -> "RenderResult":
  75
  76         stack: List[Iterator[Tuple[bool, Tree]]] = []
  77         pop = stack.pop
  78         push = stack.append
  79         new_line = Segment.line()
  80
  81         get_style = console.get_style
  82         null_style = Style.null()
  83         guide_style = get_style(self.guide_style, default="") or null_style
  84         SPACE, CONTINUE, FORK, END = range(4)
  85
  86         ASCII_GUIDES = ("    ", "|   ", "+-- ", "`-- ")
  87         TREE_GUIDES = [
  88             ("    ", "│   ", "├── ", "└── "),
  89             ("    ", "┃   ", "┣━━ ", "┗━━ "),
  90             ("    ", "║   ", "╠══ ", "╚══ "),
  91         ]
  92         _Segment = Segment
  93
  94         def make_guide(index: int, style: Style) -> Segment:
  95             """Make a Segment for a level of the guide lines."""
  96             if options.ascii_only:
  97                 line = ASCII_GUIDES[index]
  98             else:
  99                 guide = 1 if style.bold else (2 if style.underline2 else 0)
 100                 line = TREE_GUIDES[0 if options.legacy_windows else guide][index]
 101             return _Segment(line, style)
 102
 103         levels: List[Segment] = [make_guide(CONTINUE, guide_style)]
 104         push(iter(loop_last([self])))
 105
 106         guide_style_stack = StyleStack(get_style(self.guide_style))
 107         style_stack = StyleStack(get_style(self.style))
 108         remove_guide_styles = Style(bold=False, underline2=False)
 109
 110         depth = 0
 111
 112         while stack:
 113             stack_node = pop()
 114             try:
 115                 last, node = next(stack_node)
 116             except StopIteration:
 117                 levels.pop()
 118                 if levels:
 119                     guide_style = levels[-1].style or null_style
 120                     levels[-1] = make_guide(FORK, guide_style)
 121                     guide_style_stack.pop()
 122                     style_stack.pop()
 123                 continue
 124             push(stack_node)
 125             if last:
 126                 levels[-1] = make_guide(END, levels[-1].style or null_style)
 127
 128             guide_style = guide_style_stack.current + get_style(node.guide_style)
 129             style = style_stack.current + get_style(node.style)
 130             prefix = levels[(2 if self.hide_root else 1) :]
 131             renderable_lines = console.render_lines(
 132                 Styled(node.label, style),
 133                 options.update(
 134                     width=options.max_width
 135                     - sum(level.cell_length for level in prefix),
 136                     highlight=self.highlight,
 137                     height=None,
 138                 ),
 139                 pad=options.justify is not None,
 140             )
 141
 142             if not (depth == 0 and self.hide_root):
 143                 for first, line in loop_first(renderable_lines):
 144                     if prefix:
 145                         yield from _Segment.apply_style(
 146                             prefix,
 147                             style.background_style,
 148                             post_style=remove_guide_styles,
 149                         )
 150                     yield from line
 151                     yield new_line
 152                     if first and prefix:
 153                         prefix[-1] = make_guide(
 154                             SPACE if last else CONTINUE, prefix[-1].style or null_style
 155                         )
 156
 157             if node.expanded and node.children:
 158                 levels[-1] = make_guide(
 159                     SPACE if last else CONTINUE, levels[-1].style or null_style
 160                 )
 161                 levels.append(
 162                     make_guide(END if len(node.children) == 1 else FORK, guide_style)
 163                 )
 164                 style_stack.push(get_style(node.style))
 165                 guide_style_stack.push(get_style(node.guide_style))
 166                 push(iter(loop_last(node.children)))
 167                 depth += 1
 168
 169     def __rich_measure__(
 170         self, console: "Console", options: "ConsoleOptions"
 171     ) -> "Measurement":
 172         stack: List[Iterator[Tree]] = [iter([self])]
 173         pop = stack.pop
 174         push = stack.append
 175         minimum = 0
 176         maximum = 0
 177         measure = Measurement.get
 178         level = 0
 179         while stack:
 180             iter_tree = pop()
 181             try:
 182                 tree = next(iter_tree)
 183             except StopIteration:
 184                 level -= 1
 185                 continue
 186             push(iter_tree)
 187             min_measure, max_measure = measure(console, options, tree.label)
 188             indent = level * 4
 189             minimum = max(min_measure + indent, minimum)
 190             maximum = max(max_measure + indent, maximum)
 191             if tree.expanded and tree.children:
 192                 push(iter(tree.children))
 193                 level += 1
 194         return Measurement(minimum, maximum)
 195
 196
 197 if __name__ == "__main__":  # pragma: no cover
 198
 199     from pip._vendor.rich.console import Group
 200     from pip._vendor.rich.markdown import Markdown
 201     from pip._vendor.rich.panel import Panel
 202     from pip._vendor.rich.syntax import Syntax
 203     from pip._vendor.rich.table import Table
 204
 205     table = Table(row_styles=["", "dim"])
 206
 207     table.add_column("Released", style="cyan", no_wrap=True)
 208     table.add_column("Title", style="magenta")
 209     table.add_column("Box Office", justify="right", style="green")
 210
 211     table.add_row("Dec 20, 2019", "Star Wars: The Rise of Skywalker", "$952,110,690")
 212     table.add_row("May 25, 2018", "Solo: A Star Wars Story", "$393,151,347")
 213     table.add_row("Dec 15, 2017", "Star Wars Ep. V111: The Last Jedi", "$1,332,539,889")
 214     table.add_row("Dec 16, 2016", "Rogue One: A Star Wars Story", "$1,332,439,889")
 215
 216     code = """\
 217 class Segment(NamedTuple):
 218     text: str = ""
 219     style: Optional[Style] = None
 220     is_control: bool = False
 221 """
 222     syntax = Syntax(code, "python", theme="monokai", line_numbers=True)
 223
 224     markdown = Markdown(
 225         """\
 226 ### example.md
 227 > Hello, World!
 228 >
 229 > Markdown _all_ the things
 230 """
 231     )
 232
 233     root = Tree("🌲 [b green]Rich Tree", highlight=True, hide_root=True)
 234
 235     node = root.add(":file_folder: Renderables", guide_style="red")
 236     simple_node = node.add(":file_folder: [bold yellow]Atomic", guide_style="uu green")
 237     simple_node.add(Group("📄 Syntax", syntax))
 238     simple_node.add(Group("📄 Markdown", Panel(markdown, border_style="green")))
 239
 240     containers_node = node.add(
 241         ":file_folder: [bold magenta]Containers", guide_style="bold magenta"
 242     )
 243     containers_node.expanded = True
 244     panel = Panel.fit("Just a panel", border_style="red")
 245     containers_node.add(Group("📄 Panels", panel))
 246
 247     containers_node.add(Group("📄 [b magenta]Table", table))
 248
 249     console = Console()
 250
 251     console.print(root)