codekingpro/portable-devtools
115k
1from enum import IntEnum2from functools import lru_cache3from itertools import filterfalse4from operator import attrgetter5from typing import (6 TYPE_CHECKING,7 Dict,8 Iterable,9 List,10 NamedTuple,11 Optional,12 Sequence,13 Tuple,14 Type,15 Union,16)17 18from .cells import (19 _is_single_cell_widths,20 cached_cell_len,21 cell_len,22 get_character_cell_size,23 set_cell_size,24)25from .repr import Result, rich_repr26from .style import Style27 28if TYPE_CHECKING:29 from .console import Console, ConsoleOptions, RenderResult30 31 32class ControlType(IntEnum):33 """Non-printable control codes which typically translate to ANSI codes."""34 35 BELL = 136 CARRIAGE_RETURN = 237 HOME = 338 CLEAR = 439 SHOW_CURSOR = 540 HIDE_CURSOR = 641 ENABLE_ALT_SCREEN = 742 DISABLE_ALT_SCREEN = 843 CURSOR_UP = 944 CURSOR_DOWN = 1045 CURSOR_FORWARD = 1146 CURSOR_BACKWARD = 1247 CURSOR_MOVE_TO_COLUMN = 1348 CURSOR_MOVE_TO = 1449 ERASE_IN_LINE = 1550 SET_WINDOW_TITLE = 1651 52 53ControlCode = Union[54 Tuple[ControlType],55 Tuple[ControlType, Union[int, str]],56 Tuple[ControlType, int, int],57]58 59 60@rich_repr()61class Segment(NamedTuple):62 """A piece of text with associated style. Segments are produced by the Console render process and63 are ultimately converted in to strings to be written to the terminal.64 65 Args:66 text (str): A piece of text.67 style (:class:`~rich.style.Style`, optional): An optional style to apply to the text.68 control (Tuple[ControlCode], optional): Optional sequence of control codes.69 70 Attributes:71 cell_length (int): The cell length of this Segment.72 """73 74 text: str75 style: Optional[Style] = None76 control: Optional[Sequence[ControlCode]] = None77 78 @property79 def cell_length(self) -> int:80 """The number of terminal cells required to display self.text.81 82 Returns:83 int: A number of cells.84 """85 text, _style, control = self86 return 0 if control else cell_len(text)87 88 def __rich_repr__(self) -> Result:89 yield self.text90 if self.control is None:91 if self.style is not None:92 yield self.style93 else:94 yield self.style95 yield self.control96 97 def __bool__(self) -> bool:98 """Check if the segment contains text."""99 return bool(self.text)100 101 @property102 def is_control(self) -> bool:103 """Check if the segment contains control codes."""104 return self.control is not None105 106 @classmethod107 @lru_cache(1024 * 16)108 def _split_cells(cls, segment: "Segment", cut: int) -> Tuple["Segment", "Segment"]:109 """Split a segment in to two at a given cell position.110 111 Note that splitting a double-width character, may result in that character turning112 into two spaces.113 114 Args:115 segment (Segment): A segment to split.116 cut (int): A cell position to cut on.117 118 Returns:119 A tuple of two segments.120 """121 text, style, control = segment122 _Segment = Segment123 cell_length = segment.cell_length124 if cut >= cell_length:125 return segment, _Segment("", style, control)126 127 cell_size = get_character_cell_size128 129 pos = int((cut / cell_length) * len(text))130 131 while True:132 before = text[:pos]133 cell_pos = cell_len(before)134 out_by = cell_pos - cut135 if not out_by:136 return (137 _Segment(before, style, control),138 _Segment(text[pos:], style, control),139 )140 if out_by == -1 and cell_size(text[pos]) == 2:141 return (142 _Segment(text[:pos] + " ", style, control),143 _Segment(" " + text[pos + 1 :], style, control),144 )145 if out_by == +1 and cell_size(text[pos - 1]) == 2:146 return (147 _Segment(text[: pos - 1] + " ", style, control),148 _Segment(" " + text[pos:], style, control),149 )150 if cell_pos < cut:151 pos += 1152 else:153 pos -= 1154 155 def split_cells(self, cut: int) -> Tuple["Segment", "Segment"]:156 """Split segment in to two segments at the specified column.157 158 If the cut point falls in the middle of a 2-cell wide character then it is replaced159 by two spaces, to preserve the display width of the parent segment.160 161 Args:162 cut (int): Offset within the segment to cut.163 164 Returns:165 Tuple[Segment, Segment]: Two segments.166 """167 text, style, control = self168 assert cut >= 0169 170 if _is_single_cell_widths(text):171 # Fast path with all 1 cell characters172 if cut >= len(text):173 return self, Segment("", style, control)174 return (175 Segment(text[:cut], style, control),176 Segment(text[cut:], style, control),177 )178 179 return self._split_cells(self, cut)180 181 @classmethod182 def line(cls) -> "Segment":183 """Make a new line segment."""184 return cls("\n")185 186 @classmethod187 def apply_style(188 cls,189 segments: Iterable["Segment"],190 style: Optional[Style] = None,191 post_style: Optional[Style] = None,192 ) -> Iterable["Segment"]:193 """Apply style(s) to an iterable of segments.194 195 Returns an iterable of segments where the style is replaced by ``style + segment.style + post_style``.196 197 Args:198 segments (Iterable[Segment]): Segments to process.199 style (Style, optional): Base style. Defaults to None.200 post_style (Style, optional): Style to apply on top of segment style. Defaults to None.201 202 Returns:203 Iterable[Segments]: A new iterable of segments (possibly the same iterable).204 """205 result_segments = segments206 if style:207 apply = style.__add__208 result_segments = (209 cls(text, None if control else apply(_style), control)210 for text, _style, control in result_segments211 )212 if post_style:213 result_segments = (214 cls(215 text,216 (217 None218 if control219 else (_style + post_style if _style else post_style)220 ),221 control,222 )223 for text, _style, control in result_segments224 )225 return result_segments226 227 @classmethod228 def filter_control(229 cls, segments: Iterable["Segment"], is_control: bool = False230 ) -> Iterable["Segment"]:231 """Filter segments by ``is_control`` attribute.232 233 Args:234 segments (Iterable[Segment]): An iterable of Segment instances.235 is_control (bool, optional): is_control flag to match in search.236 237 Returns:238 Iterable[Segment]: And iterable of Segment instances.239 240 """241 if is_control:242 return filter(attrgetter("control"), segments)243 else:244 return filterfalse(attrgetter("control"), segments)245 246 @classmethod247 def split_lines(cls, segments: Iterable["Segment"]) -> Iterable[List["Segment"]]:248 """Split a sequence of segments in to a list of lines.249 250 Args:251 segments (Iterable[Segment]): Segments potentially containing line feeds.252 253 Yields:254 Iterable[List[Segment]]: Iterable of segment lists, one per line.255 """256 line: List[Segment] = []257 append = line.append258 259 for segment in segments:260 if "\n" in segment.text and not segment.control:261 text, style, _ = segment262 while text:263 _text, new_line, text = text.partition("\n")264 if _text:265 append(cls(_text, style))266 if new_line:267 yield line268 line = []269 append = line.append270 else:271 append(segment)272 if line:273 yield line274 275 @classmethod276 def split_lines_terminator(277 cls, segments: Iterable["Segment"]278 ) -> Iterable[Tuple[List["Segment"], bool]]:279 """Split a sequence of segments in to a list of lines and a boolean to indicate if there was a new line.280 281 Args:282 segments (Iterable[Segment]): Segments potentially containing line feeds.283 284 Yields:285 Iterable[List[Segment]]: Iterable of segment lists, one per line.286 """287 line: List[Segment] = []288 append = line.append289 290 for segment in segments:291 if "\n" in segment.text and not segment.control:292 text, style, _ = segment293 while text:294 _text, new_line, text = text.partition("\n")295 if _text:296 append(cls(_text, style))297 if new_line:298 yield (line, True)299 line = []300 append = line.append301 else:302 append(segment)303 if line:304 yield (line, False)305 306 @classmethod307 def split_and_crop_lines(308 cls,309 segments: Iterable["Segment"],310 length: int,311 style: Optional[Style] = None,312 pad: bool = True,313 include_new_lines: bool = True,314 ) -> Iterable[List["Segment"]]:315 """Split segments in to lines, and crop lines greater than a given length.316 317 Args:318 segments (Iterable[Segment]): An iterable of segments, probably319 generated from console.render.320 length (int): Desired line length.321 style (Style, optional): Style to use for any padding.322 pad (bool): Enable padding of lines that are less than `length`.323 324 Returns:325 Iterable[List[Segment]]: An iterable of lines of segments.326 """327 line: List[Segment] = []328 append = line.append329 330 adjust_line_length = cls.adjust_line_length331 new_line_segment = cls("\n")332 333 for segment in segments:334 if "\n" in segment.text and not segment.control:335 text, segment_style, _ = segment336 while text:337 _text, new_line, text = text.partition("\n")338 if _text:339 append(cls(_text, segment_style))340 if new_line:341 cropped_line = adjust_line_length(342 line, length, style=style, pad=pad343 )344 if include_new_lines:345 cropped_line.append(new_line_segment)346 yield cropped_line347 line.clear()348 else:349 append(segment)350 if line:351 yield adjust_line_length(line, length, style=style, pad=pad)352 353 @classmethod354 def adjust_line_length(355 cls,356 line: List["Segment"],357 length: int,358 style: Optional[Style] = None,359 pad: bool = True,360 ) -> List["Segment"]:361 """Adjust a line to a given width (cropping or padding as required).362 363 Args:364 segments (Iterable[Segment]): A list of segments in a single line.365 length (int): The desired width of the line.366 style (Style, optional): The style of padding if used (space on the end). Defaults to None.367 pad (bool, optional): Pad lines with spaces if they are shorter than `length`. Defaults to True.368 369 Returns:370 List[Segment]: A line of segments with the desired length.371 """372 line_length = sum(segment.cell_length for segment in line)373 new_line: List[Segment]374 375 if line_length < length:376 if pad:377 new_line = line + [cls(" " * (length - line_length), style)]378 else:379 new_line = line[:]380 elif line_length > length:381 new_line = []382 append = new_line.append383 line_length = 0384 for segment in line:385 segment_length = segment.cell_length386 if line_length + segment_length < length or segment.control:387 append(segment)388 line_length += segment_length389 else:390 text, segment_style, _ = segment391 text = set_cell_size(text, length - line_length)392 append(cls(text, segment_style))393 break394 else:395 new_line = line[:]396 return new_line397 398 @classmethod399 def get_line_length(cls, line: List["Segment"]) -> int:400 """Get the length of list of segments.401 402 Args:403 line (List[Segment]): A line encoded as a list of Segments (assumes no '\\\\n' characters),404 405 Returns:406 int: The length of the line.407 """408 _cell_len = cell_len409 return sum(_cell_len(text) for text, style, control in line if not control)410 411 @classmethod412 def get_shape(cls, lines: List[List["Segment"]]) -> Tuple[int, int]:413 """Get the shape (enclosing rectangle) of a list of lines.414 415 Args:416 lines (List[List[Segment]]): A list of lines (no '\\\\n' characters).417 418 Returns:419 Tuple[int, int]: Width and height in characters.420 """421 get_line_length = cls.get_line_length422 max_width = max(get_line_length(line) for line in lines) if lines else 0423 return (max_width, len(lines))424 425 @classmethod426 def set_shape(427 cls,428 lines: List[List["Segment"]],429 width: int,430 height: Optional[int] = None,431 style: Optional[Style] = None,432 new_lines: bool = False,433 ) -> List[List["Segment"]]:434 """Set the shape of a list of lines (enclosing rectangle).435 436 Args:437 lines (List[List[Segment]]): A list of lines.438 width (int): Desired width.439 height (int, optional): Desired height or None for no change.440 style (Style, optional): Style of any padding added.441 new_lines (bool, optional): Padded lines should include "\n". Defaults to False.442 443 Returns:444 List[List[Segment]]: New list of lines.445 """446 _height = height or len(lines)447 448 blank = (449 [cls(" " * width + "\n", style)] if new_lines else [cls(" " * width, style)]450 )451 452 adjust_line_length = cls.adjust_line_length453 shaped_lines = lines[:_height]454 shaped_lines[:] = [455 adjust_line_length(line, width, style=style) for line in lines456 ]457 if len(shaped_lines) < _height:458 shaped_lines.extend([blank] * (_height - len(shaped_lines)))459 return shaped_lines460 461 @classmethod462 def align_top(463 cls: Type["Segment"],464 lines: List[List["Segment"]],465 width: int,466 height: int,467 style: Style,468 new_lines: bool = False,469 ) -> List[List["Segment"]]:470 """Aligns lines to top (adds extra lines to bottom as required).471 472 Args:473 lines (List[List[Segment]]): A list of lines.474 width (int): Desired width.475 height (int, optional): Desired height or None for no change.476 style (Style): Style of any padding added.477 new_lines (bool, optional): Padded lines should include "\n". Defaults to False.478 479 Returns:480 List[List[Segment]]: New list of lines.481 """482 extra_lines = height - len(lines)483 if not extra_lines:484 return lines[:]485 lines = lines[:height]486 blank = cls(" " * width + "\n", style) if new_lines else cls(" " * width, style)487 lines = lines + [[blank]] * extra_lines488 return lines489 490 @classmethod491 def align_bottom(492 cls: Type["Segment"],493 lines: List[List["Segment"]],494 width: int,495 height: int,496 style: Style,497 new_lines: bool = False,498 ) -> List[List["Segment"]]:499 """Aligns render to bottom (adds extra lines above as required).500 501 Args:502 lines (List[List[Segment]]): A list of lines.503 width (int): Desired width.504 height (int, optional): Desired height or None for no change.505 style (Style): Style of any padding added. Defaults to None.506 new_lines (bool, optional): Padded lines should include "\n". Defaults to False.507 508 Returns:509 List[List[Segment]]: New list of lines.510 """511 extra_lines = height - len(lines)512 if not extra_lines:513 return lines[:]514 lines = lines[:height]515 blank = cls(" " * width + "\n", style) if new_lines else cls(" " * width, style)516 lines = [[blank]] * extra_lines + lines517 return lines518 519 @classmethod520 def align_middle(521 cls: Type["Segment"],522 lines: List[List["Segment"]],523 width: int,524 height: int,525 style: Style,526 new_lines: bool = False,527 ) -> List[List["Segment"]]:528 """Aligns lines to middle (adds extra lines to above and below as required).529 530 Args:531 lines (List[List[Segment]]): A list of lines.532 width (int): Desired width.533 height (int, optional): Desired height or None for no change.534 style (Style): Style of any padding added.535 new_lines (bool, optional): Padded lines should include "\n". Defaults to False.536 537 Returns:538 List[List[Segment]]: New list of lines.539 """540 extra_lines = height - len(lines)541 if not extra_lines:542 return lines[:]543 lines = lines[:height]544 blank = cls(" " * width + "\n", style) if new_lines else cls(" " * width, style)545 top_lines = extra_lines // 2546 bottom_lines = extra_lines - top_lines547 lines = [[blank]] * top_lines + lines + [[blank]] * bottom_lines548 return lines549 550 @classmethod551 def simplify(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:552 """Simplify an iterable of segments by combining contiguous segments with the same style.553 554 Args:555 segments (Iterable[Segment]): An iterable of segments.556 557 Returns:558 Iterable[Segment]: A possibly smaller iterable of segments that will render the same way.559 """560 iter_segments = iter(segments)561 try:562 last_segment = next(iter_segments)563 except StopIteration:564 return565 566 _Segment = Segment567 for segment in iter_segments:568 if last_segment.style == segment.style and not segment.control:569 last_segment = _Segment(570 last_segment.text + segment.text, last_segment.style571 )572 else:573 yield last_segment574 last_segment = segment575 yield last_segment576 577 @classmethod578 def strip_links(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:579 """Remove all links from an iterable of styles.580 581 Args:582 segments (Iterable[Segment]): An iterable segments.583 584 Yields:585 Segment: Segments with link removed.586 """587 for segment in segments:588 if segment.control or segment.style is None:589 yield segment590 else:591 text, style, _control = segment592 yield cls(text, style.update_link(None) if style else None)593 594 @classmethod595 def strip_styles(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:596 """Remove all styles from an iterable of segments.597 598 Args:599 segments (Iterable[Segment]): An iterable segments.600 601 Yields:602 Segment: Segments with styles replace with None603 """604 for text, _style, control in segments:605 yield cls(text, None, control)606 607 @classmethod608 def remove_color(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:609 """Remove all color from an iterable of segments.610 611 Args:612 segments (Iterable[Segment]): An iterable segments.613 614 Yields:615 Segment: Segments with colorless style.616 """617 618 cache: Dict[Style, Style] = {}619 for text, style, control in segments:620 if style:621 colorless_style = cache.get(style)622 if colorless_style is None:623 colorless_style = style.without_color624 cache[style] = colorless_style625 yield cls(text, colorless_style, control)626 else:627 yield cls(text, None, control)628 629 @classmethod630 def divide(631 cls, segments: Iterable["Segment"], cuts: Iterable[int]632 ) -> Iterable[List["Segment"]]:633 """Divides an iterable of segments in to portions.634 635 Args:636 cuts (Iterable[int]): Cell positions where to divide.637 638 Yields:639 [Iterable[List[Segment]]]: An iterable of Segments in List.640 """641 split_segments: List["Segment"] = []642 add_segment = split_segments.append643 644 iter_cuts = iter(cuts)645 646 while True:647 cut = next(iter_cuts, -1)648 if cut == -1:649 return650 if cut != 0:651 break652 yield []653 pos = 0654 655 segments_clear = split_segments.clear656 segments_copy = split_segments.copy657 658 _cell_len = cached_cell_len659 for segment in segments:660 text, _style, control = segment661 while text:662 end_pos = pos if control else pos + _cell_len(text)663 if end_pos < cut:664 add_segment(segment)665 pos = end_pos666 break667 668 if end_pos == cut:669 add_segment(segment)670 yield segments_copy()671 segments_clear()672 pos = end_pos673 674 cut = next(iter_cuts, -1)675 if cut == -1:676 if split_segments:677 yield segments_copy()678 return679 680 break681 682 else:683 before, segment = segment.split_cells(cut - pos)684 text, _style, control = segment685 add_segment(before)686 yield segments_copy()687 segments_clear()688 pos = cut689 690 cut = next(iter_cuts, -1)691 if cut == -1:692 if split_segments:693 yield segments_copy()694 return695 696 yield segments_copy()697 698 699class Segments:700 """A simple renderable to render an iterable of segments. This class may be useful if701 you want to print segments outside of a __rich_console__ method.702 703 Args:704 segments (Iterable[Segment]): An iterable of segments.705 new_lines (bool, optional): Add new lines between segments. Defaults to False.706 """707 708 def __init__(self, segments: Iterable[Segment], new_lines: bool = False) -> None:709 self.segments = list(segments)710 self.new_lines = new_lines711 712 def __rich_console__(713 self, console: "Console", options: "ConsoleOptions"714 ) -> "RenderResult":715 if self.new_lines:716 line = Segment.line()717 for segment in self.segments:718 yield segment719 yield line720 else:721 yield from self.segments722 723 724class SegmentLines:725 def __init__(self, lines: Iterable[List[Segment]], new_lines: bool = False) -> None:726 """A simple renderable containing a number of lines of segments. May be used as an intermediate727 in rendering process.728 729 Args:730 lines (Iterable[List[Segment]]): Lists of segments forming lines.731 new_lines (bool, optional): Insert new lines after each line. Defaults to False.732 """733 self.lines = list(lines)734 self.new_lines = new_lines735 736 def __rich_console__(737 self, console: "Console", options: "ConsoleOptions"738 ) -> "RenderResult":739 if self.new_lines:740 new_line = Segment.line()741 for line in self.lines:742 yield from line743 yield new_line744 else:745 for line in self.lines:746 yield from line747 748 749if __name__ == "__main__": # pragma: no cover750 from rich.console import Console751 from rich.syntax import Syntax752 from rich.text import Text753 754 code = """from rich.console import Console755console = Console()756text = Text.from_markup("Hello, [bold magenta]World[/]!")757console.print(text)"""758 759 text = Text.from_markup("Hello, [bold magenta]World[/]!")760 761 console = Console()762 763 console.rule("rich.Segment")764 console.print(765 "A Segment is the last step in the Rich render process before generating text with ANSI codes."766 )767 console.print("\nConsider the following code:\n")768 console.print(Syntax(code, "python", line_numbers=True))769 console.print()770 console.print(771 "When you call [b]print()[/b], Rich [i]renders[/i] the object in to the following:\n"772 )773 fragments = list(console.render(text))774 console.print(fragments)775 console.print()776 console.print("The Segments are then processed to produce the following output:\n")777 console.print(text)778 console.print(779 "\nYou will only need to know this if you are implementing your own Rich renderables."780 )781