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