segment.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720
  1. from enum import IntEnum
  2. from functools import lru_cache
  3. from itertools import filterfalse
  4. from logging import getLogger
  5. from operator import attrgetter
  6. from typing import (
  7. TYPE_CHECKING,
  8. Dict,
  9. Iterable,
  10. List,
  11. NamedTuple,
  12. Optional,
  13. Sequence,
  14. Tuple,
  15. Type,
  16. Union,
  17. )
  18. from .cells import (
  19. _is_single_cell_widths,
  20. cell_len,
  21. get_character_cell_size,
  22. set_cell_size,
  23. )
  24. from .repr import Result, rich_repr
  25. from .style import Style
  26. if TYPE_CHECKING:
  27. from .console import Console, ConsoleOptions, RenderResult
  28. log = getLogger("rich")
  29. class ControlType(IntEnum):
  30. """Non-printable control codes which typically translate to ANSI codes."""
  31. BELL = 1
  32. CARRIAGE_RETURN = 2
  33. HOME = 3
  34. CLEAR = 4
  35. SHOW_CURSOR = 5
  36. HIDE_CURSOR = 6
  37. ENABLE_ALT_SCREEN = 7
  38. DISABLE_ALT_SCREEN = 8
  39. CURSOR_UP = 9
  40. CURSOR_DOWN = 10
  41. CURSOR_FORWARD = 11
  42. CURSOR_BACKWARD = 12
  43. CURSOR_MOVE_TO_COLUMN = 13
  44. CURSOR_MOVE_TO = 14
  45. ERASE_IN_LINE = 15
  46. ControlCode = Union[
  47. Tuple[ControlType], Tuple[ControlType, int], Tuple[ControlType, int, int]
  48. ]
  49. @rich_repr()
  50. class Segment(NamedTuple):
  51. """A piece of text with associated style. Segments are produced by the Console render process and
  52. are ultimately converted in to strings to be written to the terminal.
  53. Args:
  54. text (str): A piece of text.
  55. style (:class:`~rich.style.Style`, optional): An optional style to apply to the text.
  56. control (Tuple[ControlCode..], optional): Optional sequence of control codes.
  57. """
  58. text: str = ""
  59. """Raw text."""
  60. style: Optional[Style] = None
  61. """An optional style."""
  62. control: Optional[Sequence[ControlCode]] = None
  63. """Optional sequence of control codes."""
  64. def __rich_repr__(self) -> Result:
  65. yield self.text
  66. if self.control is None:
  67. if self.style is not None:
  68. yield self.style
  69. else:
  70. yield self.style
  71. yield self.control
  72. def __bool__(self) -> bool:
  73. """Check if the segment contains text."""
  74. return bool(self.text)
  75. @property
  76. def cell_length(self) -> int:
  77. """Get cell length of segment."""
  78. return 0 if self.control else cell_len(self.text)
  79. @property
  80. def is_control(self) -> bool:
  81. """Check if the segment contains control codes."""
  82. return self.control is not None
  83. @classmethod
  84. @lru_cache(1024 * 16)
  85. def _split_cells(cls, segment: "Segment", cut: int) -> Tuple["Segment", "Segment"]: # type: ignore
  86. text, style, control = segment
  87. _Segment = Segment
  88. cell_length = segment.cell_length
  89. if cut >= cell_length:
  90. return segment, _Segment("", style, control)
  91. cell_size = get_character_cell_size
  92. pos = int((cut / cell_length) * len(text))
  93. before = text[:pos]
  94. cell_pos = cell_len(before)
  95. if cell_pos == cut:
  96. return (
  97. _Segment(before, style, control),
  98. _Segment(text[pos:], style, control),
  99. )
  100. while pos < len(text):
  101. char = text[pos]
  102. pos += 1
  103. cell_pos += cell_size(char)
  104. before = text[:pos]
  105. if cell_pos == cut:
  106. return (
  107. _Segment(before, style, control),
  108. _Segment(text[pos:], style, control),
  109. )
  110. if cell_pos > cut:
  111. return (
  112. _Segment(before[: pos - 1] + " ", style, control),
  113. _Segment(" " + text[pos:], style, control),
  114. )
  115. def split_cells(self, cut: int) -> Tuple["Segment", "Segment"]:
  116. """Split segment in to two segments at the specified column.
  117. If the cut point falls in the middle of a 2-cell wide character then it is replaced
  118. by two spaces, to preserve the display width of the parent segment.
  119. Returns:
  120. Tuple[Segment, Segment]: Two segments.
  121. """
  122. text, style, control = self
  123. if _is_single_cell_widths(text):
  124. # Fast path with all 1 cell characters
  125. if cut >= len(text):
  126. return self, Segment("", style, control)
  127. return (
  128. Segment(text[:cut], style, control),
  129. Segment(text[cut:], style, control),
  130. )
  131. return self._split_cells(self, cut)
  132. @classmethod
  133. def line(cls) -> "Segment":
  134. """Make a new line segment."""
  135. return cls("\n")
  136. @classmethod
  137. def apply_style(
  138. cls,
  139. segments: Iterable["Segment"],
  140. style: Optional[Style] = None,
  141. post_style: Optional[Style] = None,
  142. ) -> Iterable["Segment"]:
  143. """Apply style(s) to an iterable of segments.
  144. Returns an iterable of segments where the style is replaced by ``style + segment.style + post_style``.
  145. Args:
  146. segments (Iterable[Segment]): Segments to process.
  147. style (Style, optional): Base style. Defaults to None.
  148. post_style (Style, optional): Style to apply on top of segment style. Defaults to None.
  149. Returns:
  150. Iterable[Segments]: A new iterable of segments (possibly the same iterable).
  151. """
  152. result_segments = segments
  153. if style:
  154. apply = style.__add__
  155. result_segments = (
  156. cls(text, None if control else apply(_style), control)
  157. for text, _style, control in result_segments
  158. )
  159. if post_style:
  160. result_segments = (
  161. cls(
  162. text,
  163. (
  164. None
  165. if control
  166. else (_style + post_style if _style else post_style)
  167. ),
  168. control,
  169. )
  170. for text, _style, control in result_segments
  171. )
  172. return result_segments
  173. @classmethod
  174. def filter_control(
  175. cls, segments: Iterable["Segment"], is_control: bool = False
  176. ) -> Iterable["Segment"]:
  177. """Filter segments by ``is_control`` attribute.
  178. Args:
  179. segments (Iterable[Segment]): An iterable of Segment instances.
  180. is_control (bool, optional): is_control flag to match in search.
  181. Returns:
  182. Iterable[Segment]: And iterable of Segment instances.
  183. """
  184. if is_control:
  185. return filter(attrgetter("control"), segments)
  186. else:
  187. return filterfalse(attrgetter("control"), segments)
  188. @classmethod
  189. def split_lines(cls, segments: Iterable["Segment"]) -> Iterable[List["Segment"]]:
  190. """Split a sequence of segments in to a list of lines.
  191. Args:
  192. segments (Iterable[Segment]): Segments potentially containing line feeds.
  193. Yields:
  194. Iterable[List[Segment]]: Iterable of segment lists, one per line.
  195. """
  196. line: List[Segment] = []
  197. append = line.append
  198. for segment in segments:
  199. if "\n" in segment.text and not segment.control:
  200. text, style, _ = segment
  201. while text:
  202. _text, new_line, text = text.partition("\n")
  203. if _text:
  204. append(cls(_text, style))
  205. if new_line:
  206. yield line
  207. line = []
  208. append = line.append
  209. else:
  210. append(segment)
  211. if line:
  212. yield line
  213. @classmethod
  214. def split_and_crop_lines(
  215. cls,
  216. segments: Iterable["Segment"],
  217. length: int,
  218. style: Optional[Style] = None,
  219. pad: bool = True,
  220. include_new_lines: bool = True,
  221. ) -> Iterable[List["Segment"]]:
  222. """Split segments in to lines, and crop lines greater than a given length.
  223. Args:
  224. segments (Iterable[Segment]): An iterable of segments, probably
  225. generated from console.render.
  226. length (int): Desired line length.
  227. style (Style, optional): Style to use for any padding.
  228. pad (bool): Enable padding of lines that are less than `length`.
  229. Returns:
  230. Iterable[List[Segment]]: An iterable of lines of segments.
  231. """
  232. line: List[Segment] = []
  233. append = line.append
  234. adjust_line_length = cls.adjust_line_length
  235. new_line_segment = cls("\n")
  236. for segment in segments:
  237. if "\n" in segment.text and not segment.control:
  238. text, style, _ = segment
  239. while text:
  240. _text, new_line, text = text.partition("\n")
  241. if _text:
  242. append(cls(_text, style))
  243. if new_line:
  244. cropped_line = adjust_line_length(
  245. line, length, style=style, pad=pad
  246. )
  247. if include_new_lines:
  248. cropped_line.append(new_line_segment)
  249. yield cropped_line
  250. del line[:]
  251. else:
  252. append(segment)
  253. if line:
  254. yield adjust_line_length(line, length, style=style, pad=pad)
  255. @classmethod
  256. def adjust_line_length(
  257. cls,
  258. line: List["Segment"],
  259. length: int,
  260. style: Optional[Style] = None,
  261. pad: bool = True,
  262. ) -> List["Segment"]:
  263. """Adjust a line to a given width (cropping or padding as required).
  264. Args:
  265. segments (Iterable[Segment]): A list of segments in a single line.
  266. length (int): The desired width of the line.
  267. style (Style, optional): The style of padding if used (space on the end). Defaults to None.
  268. pad (bool, optional): Pad lines with spaces if they are shorter than `length`. Defaults to True.
  269. Returns:
  270. List[Segment]: A line of segments with the desired length.
  271. """
  272. line_length = sum(segment.cell_length for segment in line)
  273. new_line: List[Segment]
  274. if line_length < length:
  275. if pad:
  276. new_line = line + [cls(" " * (length - line_length), style)]
  277. else:
  278. new_line = line[:]
  279. elif line_length > length:
  280. new_line = []
  281. append = new_line.append
  282. line_length = 0
  283. for segment in line:
  284. segment_length = segment.cell_length
  285. if line_length + segment_length < length or segment.control:
  286. append(segment)
  287. line_length += segment_length
  288. else:
  289. text, segment_style, _ = segment
  290. text = set_cell_size(text, length - line_length)
  291. append(cls(text, segment_style))
  292. break
  293. else:
  294. new_line = line[:]
  295. return new_line
  296. @classmethod
  297. def get_line_length(cls, line: List["Segment"]) -> int:
  298. """Get the length of list of segments.
  299. Args:
  300. line (List[Segment]): A line encoded as a list of Segments (assumes no '\\\\n' characters),
  301. Returns:
  302. int: The length of the line.
  303. """
  304. _cell_len = cell_len
  305. return sum(_cell_len(segment.text) for segment in line)
  306. @classmethod
  307. def get_shape(cls, lines: List[List["Segment"]]) -> Tuple[int, int]:
  308. """Get the shape (enclosing rectangle) of a list of lines.
  309. Args:
  310. lines (List[List[Segment]]): A list of lines (no '\\\\n' characters).
  311. Returns:
  312. Tuple[int, int]: Width and height in characters.
  313. """
  314. get_line_length = cls.get_line_length
  315. max_width = max(get_line_length(line) for line in lines) if lines else 0
  316. return (max_width, len(lines))
  317. @classmethod
  318. def set_shape(
  319. cls,
  320. lines: List[List["Segment"]],
  321. width: int,
  322. height: Optional[int] = None,
  323. style: Optional[Style] = None,
  324. new_lines: bool = False,
  325. ) -> List[List["Segment"]]:
  326. """Set the shape of a list of lines (enclosing rectangle).
  327. Args:
  328. lines (List[List[Segment]]): A list of lines.
  329. width (int): Desired width.
  330. height (int, optional): Desired height or None for no change.
  331. style (Style, optional): Style of any padding added.
  332. new_lines (bool, optional): Padded lines should include "\n". Defaults to False.
  333. Returns:
  334. List[List[Segment]]: New list of lines.
  335. """
  336. _height = height or len(lines)
  337. blank = (
  338. [cls(" " * width + "\n", style)] if new_lines else [cls(" " * width, style)]
  339. )
  340. adjust_line_length = cls.adjust_line_length
  341. shaped_lines = lines[:_height]
  342. shaped_lines[:] = [
  343. adjust_line_length(line, width, style=style) for line in lines
  344. ]
  345. if len(shaped_lines) < _height:
  346. shaped_lines.extend([blank] * (_height - len(shaped_lines)))
  347. return shaped_lines
  348. @classmethod
  349. def align_top(
  350. cls: Type["Segment"],
  351. lines: List[List["Segment"]],
  352. width: int,
  353. height: int,
  354. style: Style,
  355. new_lines: bool = False,
  356. ) -> List[List["Segment"]]:
  357. """Aligns lines to top (adds extra lines to bottom as required).
  358. Args:
  359. lines (List[List[Segment]]): A list of lines.
  360. width (int): Desired width.
  361. height (int, optional): Desired height or None for no change.
  362. style (Style): Style of any padding added.
  363. new_lines (bool, optional): Padded lines should include "\n". Defaults to False.
  364. Returns:
  365. List[List[Segment]]: New list of lines.
  366. """
  367. extra_lines = height - len(lines)
  368. if not extra_lines:
  369. return lines[:]
  370. lines = lines[:height]
  371. blank = cls(" " * width + "\n", style) if new_lines else cls(" " * width, style)
  372. lines = lines + [[blank]] * extra_lines
  373. return lines
  374. @classmethod
  375. def align_bottom(
  376. cls: Type["Segment"],
  377. lines: List[List["Segment"]],
  378. width: int,
  379. height: int,
  380. style: Style,
  381. new_lines: bool = False,
  382. ) -> List[List["Segment"]]:
  383. """Aligns render to bottom (adds extra lines above as required).
  384. Args:
  385. lines (List[List[Segment]]): A list of lines.
  386. width (int): Desired width.
  387. height (int, optional): Desired height or None for no change.
  388. style (Style): Style of any padding added. Defaults to None.
  389. new_lines (bool, optional): Padded lines should include "\n". Defaults to False.
  390. Returns:
  391. List[List[Segment]]: New list of lines.
  392. """
  393. extra_lines = height - len(lines)
  394. if not extra_lines:
  395. return lines[:]
  396. lines = lines[:height]
  397. blank = cls(" " * width + "\n", style) if new_lines else cls(" " * width, style)
  398. lines = [[blank]] * extra_lines + lines
  399. return lines
  400. @classmethod
  401. def align_middle(
  402. cls: Type["Segment"],
  403. lines: List[List["Segment"]],
  404. width: int,
  405. height: int,
  406. style: Style,
  407. new_lines: bool = False,
  408. ) -> List[List["Segment"]]:
  409. """Aligns lines to middle (adds extra lines to above and below as required).
  410. Args:
  411. lines (List[List[Segment]]): A list of lines.
  412. width (int): Desired width.
  413. height (int, optional): Desired height or None for no change.
  414. style (Style): Style of any padding added.
  415. new_lines (bool, optional): Padded lines should include "\n". Defaults to False.
  416. Returns:
  417. List[List[Segment]]: New list of lines.
  418. """
  419. extra_lines = height - len(lines)
  420. if not extra_lines:
  421. return lines[:]
  422. lines = lines[:height]
  423. blank = cls(" " * width + "\n", style) if new_lines else cls(" " * width, style)
  424. top_lines = extra_lines // 2
  425. bottom_lines = extra_lines - top_lines
  426. lines = [[blank]] * top_lines + lines + [[blank]] * bottom_lines
  427. return lines
  428. @classmethod
  429. def simplify(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:
  430. """Simplify an iterable of segments by combining contiguous segments with the same style.
  431. Args:
  432. segments (Iterable[Segment]): An iterable of segments.
  433. Returns:
  434. Iterable[Segment]: A possibly smaller iterable of segments that will render the same way.
  435. """
  436. iter_segments = iter(segments)
  437. try:
  438. last_segment = next(iter_segments)
  439. except StopIteration:
  440. return
  441. _Segment = Segment
  442. for segment in iter_segments:
  443. if last_segment.style == segment.style and not segment.control:
  444. last_segment = _Segment(
  445. last_segment.text + segment.text, last_segment.style
  446. )
  447. else:
  448. yield last_segment
  449. last_segment = segment
  450. yield last_segment
  451. @classmethod
  452. def strip_links(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:
  453. """Remove all links from an iterable of styles.
  454. Args:
  455. segments (Iterable[Segment]): An iterable segments.
  456. Yields:
  457. Segment: Segments with link removed.
  458. """
  459. for segment in segments:
  460. if segment.control or segment.style is None:
  461. yield segment
  462. else:
  463. text, style, _control = segment
  464. yield cls(text, style.update_link(None) if style else None)
  465. @classmethod
  466. def strip_styles(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:
  467. """Remove all styles from an iterable of segments.
  468. Args:
  469. segments (Iterable[Segment]): An iterable segments.
  470. Yields:
  471. Segment: Segments with styles replace with None
  472. """
  473. for text, _style, control in segments:
  474. yield cls(text, None, control)
  475. @classmethod
  476. def remove_color(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:
  477. """Remove all color from an iterable of segments.
  478. Args:
  479. segments (Iterable[Segment]): An iterable segments.
  480. Yields:
  481. Segment: Segments with colorless style.
  482. """
  483. cache: Dict[Style, Style] = {}
  484. for text, style, control in segments:
  485. if style:
  486. colorless_style = cache.get(style)
  487. if colorless_style is None:
  488. colorless_style = style.without_color
  489. cache[style] = colorless_style
  490. yield cls(text, colorless_style, control)
  491. else:
  492. yield cls(text, None, control)
  493. @classmethod
  494. def divide(
  495. cls, segments: Iterable["Segment"], cuts: Iterable[int]
  496. ) -> Iterable[List["Segment"]]:
  497. """Divides an iterable of segments in to portions.
  498. Args:
  499. cuts (Iterable[int]): Cell positions where to divide.
  500. Yields:
  501. [Iterable[List[Segment]]]: An iterable of Segments in List.
  502. """
  503. split_segments: List["Segment"] = []
  504. add_segment = split_segments.append
  505. iter_cuts = iter(cuts)
  506. while True:
  507. try:
  508. cut = next(iter_cuts)
  509. except StopIteration:
  510. return []
  511. if cut != 0:
  512. break
  513. yield []
  514. pos = 0
  515. for segment in segments:
  516. while segment.text:
  517. end_pos = pos + segment.cell_length
  518. if end_pos < cut:
  519. add_segment(segment)
  520. pos = end_pos
  521. break
  522. try:
  523. if end_pos == cut:
  524. add_segment(segment)
  525. yield split_segments[:]
  526. del split_segments[:]
  527. pos = end_pos
  528. break
  529. else:
  530. before, segment = segment.split_cells(cut - pos)
  531. add_segment(before)
  532. yield split_segments[:]
  533. del split_segments[:]
  534. pos = cut
  535. finally:
  536. try:
  537. cut = next(iter_cuts)
  538. except StopIteration:
  539. if split_segments:
  540. yield split_segments[:]
  541. return
  542. yield split_segments[:]
  543. class Segments:
  544. """A simple renderable to render an iterable of segments. This class may be useful if
  545. you want to print segments outside of a __rich_console__ method.
  546. Args:
  547. segments (Iterable[Segment]): An iterable of segments.
  548. new_lines (bool, optional): Add new lines between segments. Defaults to False.
  549. """
  550. def __init__(self, segments: Iterable[Segment], new_lines: bool = False) -> None:
  551. self.segments = list(segments)
  552. self.new_lines = new_lines
  553. def __rich_console__(
  554. self, console: "Console", options: "ConsoleOptions"
  555. ) -> "RenderResult":
  556. if self.new_lines:
  557. line = Segment.line()
  558. for segment in self.segments:
  559. yield segment
  560. yield line
  561. else:
  562. yield from self.segments
  563. class SegmentLines:
  564. def __init__(self, lines: Iterable[List[Segment]], new_lines: bool = False) -> None:
  565. """A simple renderable containing a number of lines of segments. May be used as an intermediate
  566. in rendering process.
  567. Args:
  568. lines (Iterable[List[Segment]]): Lists of segments forming lines.
  569. new_lines (bool, optional): Insert new lines after each line. Defaults to False.
  570. """
  571. self.lines = list(lines)
  572. self.new_lines = new_lines
  573. def __rich_console__(
  574. self, console: "Console", options: "ConsoleOptions"
  575. ) -> "RenderResult":
  576. if self.new_lines:
  577. new_line = Segment.line()
  578. for line in self.lines:
  579. yield from line
  580. yield new_line
  581. else:
  582. for line in self.lines:
  583. yield from line
  584. if __name__ == "__main__":
  585. if __name__ == "__main__": # pragma: no cover
  586. from pip._vendor.rich.console import Console
  587. from pip._vendor.rich.syntax import Syntax
  588. from pip._vendor.rich.text import Text
  589. code = """from rich.console import Console
  590. console = Console()
  591. text = Text.from_markup("Hello, [bold magenta]World[/]!")
  592. console.print(text)"""
  593. text = Text.from_markup("Hello, [bold magenta]World[/]!")
  594. console = Console()
  595. console.rule("rich.Segment")
  596. console.print(
  597. "A Segment is the last step in the Rich render process before generating text with ANSI codes."
  598. )
  599. console.print("\nConsider the following code:\n")
  600. console.print(Syntax(code, "python", line_numbers=True))
  601. console.print()
  602. console.print(
  603. "When you call [b]print()[/b], Rich [i]renders[/i] the object in to the the following:\n"
  604. )
  605. fragments = list(console.render(text))
  606. console.print(fragments)
  607. console.print()
  608. console.print(
  609. "The Segments are then processed to produce the following output:\n"
  610. )
  611. console.print(text)
  612. console.print(
  613. "\nYou will only need to know this if you are implementing your own Rich renderables."
  614. )