tree.py 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249
  1. from typing import Iterator, List, Optional, Tuple
  2. from ._loop import loop_first, loop_last
  3. from .console import Console, ConsoleOptions, RenderableType, RenderResult
  4. from .jupyter import JupyterMixin
  5. from .measure import Measurement
  6. from .segment import Segment
  7. from .style import Style, StyleStack, StyleType
  8. from .styled import Styled
  9. class Tree(JupyterMixin):
  10. """A renderable for a tree structure.
  11. Args:
  12. label (RenderableType): The renderable or str for the tree label.
  13. style (StyleType, optional): Style of this tree. Defaults to "tree".
  14. guide_style (StyleType, optional): Style of the guide lines. Defaults to "tree.line".
  15. expanded (bool, optional): Also display children. Defaults to True.
  16. highlight (bool, optional): Highlight renderable (if str). Defaults to False.
  17. """
  18. def __init__(
  19. self,
  20. label: RenderableType,
  21. *,
  22. style: StyleType = "tree",
  23. guide_style: StyleType = "tree.line",
  24. expanded: bool = True,
  25. highlight: bool = False,
  26. hide_root: bool = False,
  27. ) -> None:
  28. self.label = label
  29. self.style = style
  30. self.guide_style = guide_style
  31. self.children: List[Tree] = []
  32. self.expanded = expanded
  33. self.highlight = highlight
  34. self.hide_root = hide_root
  35. def add(
  36. self,
  37. label: RenderableType,
  38. *,
  39. style: Optional[StyleType] = None,
  40. guide_style: Optional[StyleType] = None,
  41. expanded: bool = True,
  42. highlight: bool = False,
  43. ) -> "Tree":
  44. """Add a child tree.
  45. Args:
  46. label (RenderableType): The renderable or str for the tree label.
  47. style (StyleType, optional): Style of this tree. Defaults to "tree".
  48. guide_style (StyleType, optional): Style of the guide lines. Defaults to "tree.line".
  49. expanded (bool, optional): Also display children. Defaults to True.
  50. highlight (Optional[bool], optional): Highlight renderable (if str). Defaults to False.
  51. Returns:
  52. Tree: A new child Tree, which may be further modified.
  53. """
  54. node = Tree(
  55. label,
  56. style=self.style if style is None else style,
  57. guide_style=self.guide_style if guide_style is None else guide_style,
  58. expanded=expanded,
  59. highlight=self.highlight if highlight is None else highlight,
  60. )
  61. self.children.append(node)
  62. return node
  63. def __rich_console__(
  64. self, console: "Console", options: "ConsoleOptions"
  65. ) -> "RenderResult":
  66. stack: List[Iterator[Tuple[bool, Tree]]] = []
  67. pop = stack.pop
  68. push = stack.append
  69. new_line = Segment.line()
  70. get_style = console.get_style
  71. null_style = Style.null()
  72. guide_style = get_style(self.guide_style, default="") or null_style
  73. SPACE, CONTINUE, FORK, END = range(4)
  74. ASCII_GUIDES = (" ", "| ", "+-- ", "`-- ")
  75. TREE_GUIDES = [
  76. (" ", "│ ", "├── ", "└── "),
  77. (" ", "┃ ", "┣━━ ", "┗━━ "),
  78. (" ", "║ ", "╠══ ", "╚══ "),
  79. ]
  80. _Segment = Segment
  81. def make_guide(index: int, style: Style) -> Segment:
  82. """Make a Segment for a level of the guide lines."""
  83. if options.ascii_only:
  84. line = ASCII_GUIDES[index]
  85. else:
  86. guide = 1 if style.bold else (2 if style.underline2 else 0)
  87. line = TREE_GUIDES[0 if options.legacy_windows else guide][index]
  88. return _Segment(line, style)
  89. levels: List[Segment] = [make_guide(CONTINUE, guide_style)]
  90. push(iter(loop_last([self])))
  91. guide_style_stack = StyleStack(get_style(self.guide_style))
  92. style_stack = StyleStack(get_style(self.style))
  93. remove_guide_styles = Style(bold=False, underline2=False)
  94. depth = 0
  95. while stack:
  96. stack_node = pop()
  97. try:
  98. last, node = next(stack_node)
  99. except StopIteration:
  100. levels.pop()
  101. if levels:
  102. guide_style = levels[-1].style or null_style
  103. levels[-1] = make_guide(FORK, guide_style)
  104. guide_style_stack.pop()
  105. style_stack.pop()
  106. continue
  107. push(stack_node)
  108. if last:
  109. levels[-1] = make_guide(END, levels[-1].style or null_style)
  110. guide_style = guide_style_stack.current + get_style(node.guide_style)
  111. style = style_stack.current + get_style(node.style)
  112. prefix = levels[(2 if self.hide_root else 1) :]
  113. renderable_lines = console.render_lines(
  114. Styled(node.label, style),
  115. options.update(
  116. width=options.max_width
  117. - sum(level.cell_length for level in prefix),
  118. highlight=self.highlight,
  119. height=None,
  120. ),
  121. )
  122. if not (depth == 0 and self.hide_root):
  123. for first, line in loop_first(renderable_lines):
  124. if prefix:
  125. yield from _Segment.apply_style(
  126. prefix,
  127. style.background_style,
  128. post_style=remove_guide_styles,
  129. )
  130. yield from line
  131. yield new_line
  132. if first and prefix:
  133. prefix[-1] = make_guide(
  134. SPACE if last else CONTINUE, prefix[-1].style or null_style
  135. )
  136. if node.expanded and node.children:
  137. levels[-1] = make_guide(
  138. SPACE if last else CONTINUE, levels[-1].style or null_style
  139. )
  140. levels.append(
  141. make_guide(END if len(node.children) == 1 else FORK, guide_style)
  142. )
  143. style_stack.push(get_style(node.style))
  144. guide_style_stack.push(get_style(node.guide_style))
  145. push(iter(loop_last(node.children)))
  146. depth += 1
  147. def __rich_measure__(
  148. self, console: "Console", options: "ConsoleOptions"
  149. ) -> "Measurement":
  150. stack: List[Iterator[Tree]] = [iter([self])]
  151. pop = stack.pop
  152. push = stack.append
  153. minimum = 0
  154. maximum = 0
  155. measure = Measurement.get
  156. level = 0
  157. while stack:
  158. iter_tree = pop()
  159. try:
  160. tree = next(iter_tree)
  161. except StopIteration:
  162. level -= 1
  163. continue
  164. push(iter_tree)
  165. min_measure, max_measure = measure(console, options, tree.label)
  166. indent = level * 4
  167. minimum = max(min_measure + indent, minimum)
  168. maximum = max(max_measure + indent, maximum)
  169. if tree.expanded and tree.children:
  170. push(iter(tree.children))
  171. level += 1
  172. return Measurement(minimum, maximum)
  173. if __name__ == "__main__": # pragma: no cover
  174. from pip._vendor.rich.console import Group
  175. from pip._vendor.rich.markdown import Markdown
  176. from pip._vendor.rich.panel import Panel
  177. from pip._vendor.rich.syntax import Syntax
  178. from pip._vendor.rich.table import Table
  179. table = Table(row_styles=["", "dim"])
  180. table.add_column("Released", style="cyan", no_wrap=True)
  181. table.add_column("Title", style="magenta")
  182. table.add_column("Box Office", justify="right", style="green")
  183. table.add_row("Dec 20, 2019", "Star Wars: The Rise of Skywalker", "$952,110,690")
  184. table.add_row("May 25, 2018", "Solo: A Star Wars Story", "$393,151,347")
  185. table.add_row("Dec 15, 2017", "Star Wars Ep. V111: The Last Jedi", "$1,332,539,889")
  186. table.add_row("Dec 16, 2016", "Rogue One: A Star Wars Story", "$1,332,439,889")
  187. code = """\
  188. class Segment(NamedTuple):
  189. text: str = ""
  190. style: Optional[Style] = None
  191. is_control: bool = False
  192. """
  193. syntax = Syntax(code, "python", theme="monokai", line_numbers=True)
  194. markdown = Markdown(
  195. """\
  196. ### example.md
  197. > Hello, World!
  198. >
  199. > Markdown _all_ the things
  200. """
  201. )
  202. root = Tree("🌲 [b green]Rich Tree", highlight=True, hide_root=True)
  203. node = root.add(":file_folder: Renderables", guide_style="red")
  204. simple_node = node.add(":file_folder: [bold yellow]Atomic", guide_style="uu green")
  205. simple_node.add(Group("📄 Syntax", syntax))
  206. simple_node.add(Group("📄 Markdown", Panel(markdown, border_style="green")))
  207. containers_node = node.add(
  208. ":file_folder: [bold magenta]Containers", guide_style="bold magenta"
  209. )
  210. containers_node.expanded = True
  211. panel = Panel.fit("Just a panel", border_style="red")
  212. containers_node.add(Group("📄 Panels", panel))
  213. containers_node.add(Group("📄 [b magenta]Table", table))
  214. console = Console()
  215. console.print(root)