CoolFace
Apppublic

Aluode/PerceptionLabPortable

sourceHugging Faceupdated 9mo agoView on Hugging Face
0likes
segment.py753 linesDownload Raw Back to rich
1from enum import IntEnum2from functools import lru_cache3from itertools import filterfalse4from logging import getLogger5from operator import attrgetter6from 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_repr27from .style import Style28 29if TYPE_CHECKING:30    from .console import Console, ConsoleOptions, RenderResult31 32log = getLogger("rich")33 34 35class ControlType(IntEnum):36    """Non-printable control codes which typically translate to ANSI codes."""37 38    BELL = 139    CARRIAGE_RETURN = 240    HOME = 341    CLEAR = 442    SHOW_CURSOR = 543    HIDE_CURSOR = 644    ENABLE_ALT_SCREEN = 745    DISABLE_ALT_SCREEN = 846    CURSOR_UP = 947    CURSOR_DOWN = 1048    CURSOR_FORWARD = 1149    CURSOR_BACKWARD = 1250    CURSOR_MOVE_TO_COLUMN = 1351    CURSOR_MOVE_TO = 1452    ERASE_IN_LINE = 1553    SET_WINDOW_TITLE = 1654 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 and66    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: str78    style: Optional[Style] = None79    control: Optional[Sequence[ControlCode]] = None80 81    @property82    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 = self89        return 0 if control else cell_len(text)90 91    def __rich_repr__(self) -> Result:92        yield self.text93        if self.control is None:94            if self.style is not None:95                yield self.style96        else:97            yield self.style98            yield self.control99 100    def __bool__(self) -> bool:101        """Check if the segment contains text."""102        return bool(self.text)103 104    @property105    def is_control(self) -> bool:106        """Check if the segment contains control codes."""107        return self.control is not None108 109    @classmethod110    @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 turning115        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 = segment125        _Segment = Segment126        cell_length = segment.cell_length127        if cut >= cell_length:128            return segment, _Segment("", style, control)129 130        cell_size = get_character_cell_size131 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 - cut138            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 += 1155            else:156                pos -= 1157 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 replaced162        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 = self171        assert cut >= 0172 173        if _is_single_cell_widths(text):174            # Fast path with all 1 cell characters175            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    @classmethod185    def line(cls) -> "Segment":186        """Make a new line segment."""187        return cls("\n")188 189    @classmethod190    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 = segments209        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_segments214            )215        if post_style:216            result_segments = (217                cls(218                    text,219                    (220                        None221                        if control222                        else (_style + post_style if _style else post_style)223                    ),224                    control,225                )226                for text, _style, control in result_segments227            )228        return result_segments229 230    @classmethod231    def filter_control(232        cls, segments: Iterable["Segment"], is_control: bool = False233    ) -> 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    @classmethod250    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.append261 262        for segment in segments:263            if "\n" in segment.text and not segment.control:264                text, style, _ = segment265                while text:266                    _text, new_line, text = text.partition("\n")267                    if _text:268                        append(cls(_text, style))269                    if new_line:270                        yield line271                        line = []272                        append = line.append273            else:274                append(segment)275        if line:276            yield line277 278    @classmethod279    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, probably291                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.append301 302        adjust_line_length = cls.adjust_line_length303        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, _ = segment308                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=pad315                        )316                        if include_new_lines:317                            cropped_line.append(new_line_segment)318                        yield cropped_line319                        line.clear()320            else:321                append(segment)322        if line:323            yield adjust_line_length(line, length, style=style, pad=pad)324 325    @classmethod326    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.append355            line_length = 0356            for segment in line:357                segment_length = segment.cell_length358                if line_length + segment_length < length or segment.control:359                    append(segment)360                    line_length += segment_length361                else:362                    text, segment_style, _ = segment363                    text = set_cell_size(text, length - line_length)364                    append(cls(text, segment_style))365                    break366        else:367            new_line = line[:]368        return new_line369 370    @classmethod371    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_len381        return sum(_cell_len(text) for text, style, control in line if not control)382 383    @classmethod384    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_length394        max_width = max(get_line_length(line) for line in lines) if lines else 0395        return (max_width, len(lines))396 397    @classmethod398    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_length425        shaped_lines = lines[:_height]426        shaped_lines[:] = [427            adjust_line_length(line, width, style=style) for line in lines428        ]429        if len(shaped_lines) < _height:430            shaped_lines.extend([blank] * (_height - len(shaped_lines)))431        return shaped_lines432 433    @classmethod434    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_lines460        return lines461 462    @classmethod463    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 + lines489        return lines490 491    @classmethod492    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 // 2518        bottom_lines = extra_lines - top_lines519        lines = [[blank]] * top_lines + lines + [[blank]] * bottom_lines520        return lines521 522    @classmethod523    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            return537 538        _Segment = Segment539        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.style543                )544            else:545                yield last_segment546                last_segment = segment547        yield last_segment548 549    @classmethod550    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 segment562            else:563                text, style, _control = segment564                yield cls(text, style.update_link(None) if style else None)565 566    @classmethod567    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 None575        """576        for text, _style, control in segments:577            yield cls(text, None, control)578 579    @classmethod580    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_color596                    cache[style] = colorless_style597                yield cls(text, colorless_style, control)598            else:599                yield cls(text, None, control)600 601    @classmethod602    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.append615 616        iter_cuts = iter(cuts)617 618        while True:619            cut = next(iter_cuts, -1)620            if cut == -1:621                return622            if cut != 0:623                break624            yield []625        pos = 0626 627        segments_clear = split_segments.clear628        segments_copy = split_segments.copy629 630        _cell_len = cached_cell_len631        for segment in segments:632            text, _style, control = segment633            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_pos638                    break639 640                if end_pos == cut:641                    add_segment(segment)642                    yield segments_copy()643                    segments_clear()644                    pos = end_pos645 646                    cut = next(iter_cuts, -1)647                    if cut == -1:648                        if split_segments:649                            yield segments_copy()650                        return651 652                    break653 654                else:655                    before, segment = segment.split_cells(cut - pos)656                    text, _style, control = segment657                    add_segment(before)658                    yield segments_copy()659                    segments_clear()660                    pos = cut661 662                cut = next(iter_cuts, -1)663                if cut == -1:664                    if split_segments:665                        yield segments_copy()666                    return667 668        yield segments_copy()669 670 671class Segments:672    """A simple renderable to render an iterable of segments. This class may be useful if673    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_lines683 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 segment691                yield line692        else:693            yield from self.segments694 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 intermediate699        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_lines707 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 line715                yield new_line716        else:717            for line in self.lines:718                yield from line719 720 721if __name__ == "__main__":  # pragma: no cover722    from pip._vendor.rich.console import Console723    from pip._vendor.rich.syntax import Syntax724    from pip._vendor.rich.text import Text725 726    code = """from rich.console import Console727console = 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    )753 
Aluode/PerceptionLabPortable · CoolFace