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