python (3.11.7)
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 ESC[4;38;5;81mTree(ESC[4;38;5;149mJupyterMixin):
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)
Browse
Usage
Versions