Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes14kdownloads
segment.py739 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        text, style, control = segment113        _Segment = Segment114 115        cell_length = segment.cell_length116        if cut >= cell_length:117            return segment, _Segment("", style, control)118 119        cell_size = get_character_cell_size120 121        pos = int((cut / cell_length) * (len(text) - 1))122 123        before = text[:pos]124        cell_pos = cell_len(before)125        if cell_pos == cut:126            return (127                _Segment(before, style, control),128                _Segment(text[pos:], style, control),129            )130        while pos < len(text):131            char = text[pos]132            pos += 1133            cell_pos += cell_size(char)134            before = text[:pos]135            if cell_pos == cut:136                return (137                    _Segment(before, style, control),138                    _Segment(text[pos:], style, control),139                )140            if cell_pos > cut:141                return (142                    _Segment(before[: pos - 1] + " ", style, control),143                    _Segment(" " + text[pos:], style, control),144                )145 146        raise AssertionError("Will never reach here")147 148    def split_cells(self, cut: int) -> Tuple["Segment", "Segment"]:149        """Split segment in to two segments at the specified column.150 151        If the cut point falls in the middle of a 2-cell wide character then it is replaced152        by two spaces, to preserve the display width of the parent segment.153 154        Returns:155            Tuple[Segment, Segment]: Two segments.156        """157        text, style, control = self158 159        if _is_single_cell_widths(text):160            # Fast path with all 1 cell characters161            if cut >= len(text):162                return self, Segment("", style, control)163            return (164                Segment(text[:cut], style, control),165                Segment(text[cut:], style, control),166            )167 168        return self._split_cells(self, cut)169 170    @classmethod171    def line(cls) -> "Segment":172        """Make a new line segment."""173        return cls("\n")174 175    @classmethod176    def apply_style(177        cls,178        segments: Iterable["Segment"],179        style: Optional[Style] = None,180        post_style: Optional[Style] = None,181    ) -> Iterable["Segment"]:182        """Apply style(s) to an iterable of segments.183 184        Returns an iterable of segments where the style is replaced by ``style + segment.style + post_style``.185 186        Args:187            segments (Iterable[Segment]): Segments to process.188            style (Style, optional): Base style. Defaults to None.189            post_style (Style, optional): Style to apply on top of segment style. Defaults to None.190 191        Returns:192            Iterable[Segments]: A new iterable of segments (possibly the same iterable).193        """194        result_segments = segments195        if style:196            apply = style.__add__197            result_segments = (198                cls(text, None if control else apply(_style), control)199                for text, _style, control in result_segments200            )201        if post_style:202            result_segments = (203                cls(204                    text,205                    (206                        None207                        if control208                        else (_style + post_style if _style else post_style)209                    ),210                    control,211                )212                for text, _style, control in result_segments213            )214        return result_segments215 216    @classmethod217    def filter_control(218        cls, segments: Iterable["Segment"], is_control: bool = False219    ) -> Iterable["Segment"]:220        """Filter segments by ``is_control`` attribute.221 222        Args:223            segments (Iterable[Segment]): An iterable of Segment instances.224            is_control (bool, optional): is_control flag to match in search.225 226        Returns:227            Iterable[Segment]: And iterable of Segment instances.228 229        """230        if is_control:231            return filter(attrgetter("control"), segments)232        else:233            return filterfalse(attrgetter("control"), segments)234 235    @classmethod236    def split_lines(cls, segments: Iterable["Segment"]) -> Iterable[List["Segment"]]:237        """Split a sequence of segments in to a list of lines.238 239        Args:240            segments (Iterable[Segment]): Segments potentially containing line feeds.241 242        Yields:243            Iterable[List[Segment]]: Iterable of segment lists, one per line.244        """245        line: List[Segment] = []246        append = line.append247 248        for segment in segments:249            if "\n" in segment.text and not segment.control:250                text, style, _ = segment251                while text:252                    _text, new_line, text = text.partition("\n")253                    if _text:254                        append(cls(_text, style))255                    if new_line:256                        yield line257                        line = []258                        append = line.append259            else:260                append(segment)261        if line:262            yield line263 264    @classmethod265    def split_and_crop_lines(266        cls,267        segments: Iterable["Segment"],268        length: int,269        style: Optional[Style] = None,270        pad: bool = True,271        include_new_lines: bool = True,272    ) -> Iterable[List["Segment"]]:273        """Split segments in to lines, and crop lines greater than a given length.274 275        Args:276            segments (Iterable[Segment]): An iterable of segments, probably277                generated from console.render.278            length (int): Desired line length.279            style (Style, optional): Style to use for any padding.280            pad (bool): Enable padding of lines that are less than `length`.281 282        Returns:283            Iterable[List[Segment]]: An iterable of lines of segments.284        """285        line: List[Segment] = []286        append = line.append287 288        adjust_line_length = cls.adjust_line_length289        new_line_segment = cls("\n")290 291        for segment in segments:292            if "\n" in segment.text and not segment.control:293                text, segment_style, _ = segment294                while text:295                    _text, new_line, text = text.partition("\n")296                    if _text:297                        append(cls(_text, segment_style))298                    if new_line:299                        cropped_line = adjust_line_length(300                            line, length, style=style, pad=pad301                        )302                        if include_new_lines:303                            cropped_line.append(new_line_segment)304                        yield cropped_line305                        line.clear()306            else:307                append(segment)308        if line:309            yield adjust_line_length(line, length, style=style, pad=pad)310 311    @classmethod312    def adjust_line_length(313        cls,314        line: List["Segment"],315        length: int,316        style: Optional[Style] = None,317        pad: bool = True,318    ) -> List["Segment"]:319        """Adjust a line to a given width (cropping or padding as required).320 321        Args:322            segments (Iterable[Segment]): A list of segments in a single line.323            length (int): The desired width of the line.324            style (Style, optional): The style of padding if used (space on the end). Defaults to None.325            pad (bool, optional): Pad lines with spaces if they are shorter than `length`. Defaults to True.326 327        Returns:328            List[Segment]: A line of segments with the desired length.329        """330        line_length = sum(segment.cell_length for segment in line)331        new_line: List[Segment]332 333        if line_length < length:334            if pad:335                new_line = line + [cls(" " * (length - line_length), style)]336            else:337                new_line = line[:]338        elif line_length > length:339            new_line = []340            append = new_line.append341            line_length = 0342            for segment in line:343                segment_length = segment.cell_length344                if line_length + segment_length < length or segment.control:345                    append(segment)346                    line_length += segment_length347                else:348                    text, segment_style, _ = segment349                    text = set_cell_size(text, length - line_length)350                    append(cls(text, segment_style))351                    break352        else:353            new_line = line[:]354        return new_line355 356    @classmethod357    def get_line_length(cls, line: List["Segment"]) -> int:358        """Get the length of list of segments.359 360        Args:361            line (List[Segment]): A line encoded as a list of Segments (assumes no '\\\\n' characters),362 363        Returns:364            int: The length of the line.365        """366        _cell_len = cell_len367        return sum(_cell_len(text) for text, style, control in line if not control)368 369    @classmethod370    def get_shape(cls, lines: List[List["Segment"]]) -> Tuple[int, int]:371        """Get the shape (enclosing rectangle) of a list of lines.372 373        Args:374            lines (List[List[Segment]]): A list of lines (no '\\\\n' characters).375 376        Returns:377            Tuple[int, int]: Width and height in characters.378        """379        get_line_length = cls.get_line_length380        max_width = max(get_line_length(line) for line in lines) if lines else 0381        return (max_width, len(lines))382 383    @classmethod384    def set_shape(385        cls,386        lines: List[List["Segment"]],387        width: int,388        height: Optional[int] = None,389        style: Optional[Style] = None,390        new_lines: bool = False,391    ) -> List[List["Segment"]]:392        """Set the shape of a list of lines (enclosing rectangle).393 394        Args:395            lines (List[List[Segment]]): A list of lines.396            width (int): Desired width.397            height (int, optional): Desired height or None for no change.398            style (Style, optional): Style of any padding added.399            new_lines (bool, optional): Padded lines should include "\n". Defaults to False.400 401        Returns:402            List[List[Segment]]: New list of lines.403        """404        _height = height or len(lines)405 406        blank = (407            [cls(" " * width + "\n", style)] if new_lines else [cls(" " * width, style)]408        )409 410        adjust_line_length = cls.adjust_line_length411        shaped_lines = lines[:_height]412        shaped_lines[:] = [413            adjust_line_length(line, width, style=style) for line in lines414        ]415        if len(shaped_lines) < _height:416            shaped_lines.extend([blank] * (_height - len(shaped_lines)))417        return shaped_lines418 419    @classmethod420    def align_top(421        cls: Type["Segment"],422        lines: List[List["Segment"]],423        width: int,424        height: int,425        style: Style,426        new_lines: bool = False,427    ) -> List[List["Segment"]]:428        """Aligns lines to top (adds extra lines to bottom as required).429 430        Args:431            lines (List[List[Segment]]): A list of lines.432            width (int): Desired width.433            height (int, optional): Desired height or None for no change.434            style (Style): Style of any padding added.435            new_lines (bool, optional): Padded lines should include "\n". Defaults to False.436 437        Returns:438            List[List[Segment]]: New list of lines.439        """440        extra_lines = height - len(lines)441        if not extra_lines:442            return lines[:]443        lines = lines[:height]444        blank = cls(" " * width + "\n", style) if new_lines else cls(" " * width, style)445        lines = lines + [[blank]] * extra_lines446        return lines447 448    @classmethod449    def align_bottom(450        cls: Type["Segment"],451        lines: List[List["Segment"]],452        width: int,453        height: int,454        style: Style,455        new_lines: bool = False,456    ) -> List[List["Segment"]]:457        """Aligns render to bottom (adds extra lines above as required).458 459        Args:460            lines (List[List[Segment]]): A list of lines.461            width (int): Desired width.462            height (int, optional): Desired height or None for no change.463            style (Style): Style of any padding added. Defaults to None.464            new_lines (bool, optional): Padded lines should include "\n". Defaults to False.465 466        Returns:467            List[List[Segment]]: New list of lines.468        """469        extra_lines = height - len(lines)470        if not extra_lines:471            return lines[:]472        lines = lines[:height]473        blank = cls(" " * width + "\n", style) if new_lines else cls(" " * width, style)474        lines = [[blank]] * extra_lines + lines475        return lines476 477    @classmethod478    def align_middle(479        cls: Type["Segment"],480        lines: List[List["Segment"]],481        width: int,482        height: int,483        style: Style,484        new_lines: bool = False,485    ) -> List[List["Segment"]]:486        """Aligns lines to middle (adds extra lines to above and below as required).487 488        Args:489            lines (List[List[Segment]]): A list of lines.490            width (int): Desired width.491            height (int, optional): Desired height or None for no change.492            style (Style): Style of any padding added.493            new_lines (bool, optional): Padded lines should include "\n". Defaults to False.494 495        Returns:496            List[List[Segment]]: New list of lines.497        """498        extra_lines = height - len(lines)499        if not extra_lines:500            return lines[:]501        lines = lines[:height]502        blank = cls(" " * width + "\n", style) if new_lines else cls(" " * width, style)503        top_lines = extra_lines // 2504        bottom_lines = extra_lines - top_lines505        lines = [[blank]] * top_lines + lines + [[blank]] * bottom_lines506        return lines507 508    @classmethod509    def simplify(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:510        """Simplify an iterable of segments by combining contiguous segments with the same style.511 512        Args:513            segments (Iterable[Segment]): An iterable of segments.514 515        Returns:516            Iterable[Segment]: A possibly smaller iterable of segments that will render the same way.517        """518        iter_segments = iter(segments)519        try:520            last_segment = next(iter_segments)521        except StopIteration:522            return523 524        _Segment = Segment525        for segment in iter_segments:526            if last_segment.style == segment.style and not segment.control:527                last_segment = _Segment(528                    last_segment.text + segment.text, last_segment.style529                )530            else:531                yield last_segment532                last_segment = segment533        yield last_segment534 535    @classmethod536    def strip_links(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:537        """Remove all links from an iterable of styles.538 539        Args:540            segments (Iterable[Segment]): An iterable segments.541 542        Yields:543            Segment: Segments with link removed.544        """545        for segment in segments:546            if segment.control or segment.style is None:547                yield segment548            else:549                text, style, _control = segment550                yield cls(text, style.update_link(None) if style else None)551 552    @classmethod553    def strip_styles(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:554        """Remove all styles from an iterable of segments.555 556        Args:557            segments (Iterable[Segment]): An iterable segments.558 559        Yields:560            Segment: Segments with styles replace with None561        """562        for text, _style, control in segments:563            yield cls(text, None, control)564 565    @classmethod566    def remove_color(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:567        """Remove all color from an iterable of segments.568 569        Args:570            segments (Iterable[Segment]): An iterable segments.571 572        Yields:573            Segment: Segments with colorless style.574        """575 576        cache: Dict[Style, Style] = {}577        for text, style, control in segments:578            if style:579                colorless_style = cache.get(style)580                if colorless_style is None:581                    colorless_style = style.without_color582                    cache[style] = colorless_style583                yield cls(text, colorless_style, control)584            else:585                yield cls(text, None, control)586 587    @classmethod588    def divide(589        cls, segments: Iterable["Segment"], cuts: Iterable[int]590    ) -> Iterable[List["Segment"]]:591        """Divides an iterable of segments in to portions.592 593        Args:594            cuts (Iterable[int]): Cell positions where to divide.595 596        Yields:597            [Iterable[List[Segment]]]: An iterable of Segments in List.598        """599        split_segments: List["Segment"] = []600        add_segment = split_segments.append601 602        iter_cuts = iter(cuts)603 604        while True:605            cut = next(iter_cuts, -1)606            if cut == -1:607                return []608            if cut != 0:609                break610            yield []611        pos = 0612 613        segments_clear = split_segments.clear614        segments_copy = split_segments.copy615 616        _cell_len = cached_cell_len617        for segment in segments:618            text, _style, control = segment619            while text:620                end_pos = pos if control else pos + _cell_len(text)621                if end_pos < cut:622                    add_segment(segment)623                    pos = end_pos624                    break625 626                if end_pos == cut:627                    add_segment(segment)628                    yield segments_copy()629                    segments_clear()630                    pos = end_pos631 632                    cut = next(iter_cuts, -1)633                    if cut == -1:634                        if split_segments:635                            yield segments_copy()636                        return637 638                    break639 640                else:641                    before, segment = segment.split_cells(cut - pos)642                    text, _style, control = segment643                    add_segment(before)644                    yield segments_copy()645                    segments_clear()646                    pos = cut647 648                cut = next(iter_cuts, -1)649                if cut == -1:650                    if split_segments:651                        yield segments_copy()652                    return653 654        yield segments_copy()655 656 657class Segments:658    """A simple renderable to render an iterable of segments. This class may be useful if659    you want to print segments outside of a __rich_console__ method.660 661    Args:662        segments (Iterable[Segment]): An iterable of segments.663        new_lines (bool, optional): Add new lines between segments. Defaults to False.664    """665 666    def __init__(self, segments: Iterable[Segment], new_lines: bool = False) -> None:667        self.segments = list(segments)668        self.new_lines = new_lines669 670    def __rich_console__(671        self, console: "Console", options: "ConsoleOptions"672    ) -> "RenderResult":673        if self.new_lines:674            line = Segment.line()675            for segment in self.segments:676                yield segment677                yield line678        else:679            yield from self.segments680 681 682class SegmentLines:683    def __init__(self, lines: Iterable[List[Segment]], new_lines: bool = False) -> None:684        """A simple renderable containing a number of lines of segments. May be used as an intermediate685        in rendering process.686 687        Args:688            lines (Iterable[List[Segment]]): Lists of segments forming lines.689            new_lines (bool, optional): Insert new lines after each line. Defaults to False.690        """691        self.lines = list(lines)692        self.new_lines = new_lines693 694    def __rich_console__(695        self, console: "Console", options: "ConsoleOptions"696    ) -> "RenderResult":697        if self.new_lines:698            new_line = Segment.line()699            for line in self.lines:700                yield from line701                yield new_line702        else:703            for line in self.lines:704                yield from line705 706 707if __name__ == "__main__":  # pragma: no cover708    from rich.console import Console709    from rich.syntax import Syntax710    from rich.text import Text711 712    code = """from rich.console import Console713console = Console()714text = Text.from_markup("Hello, [bold magenta]World[/]!")715console.print(text)"""716 717    text = Text.from_markup("Hello, [bold magenta]World[/]!")718 719    console = Console()720 721    console.rule("rich.Segment")722    console.print(723        "A Segment is the last step in the Rich render process before generating text with ANSI codes."724    )725    console.print("\nConsider the following code:\n")726    console.print(Syntax(code, "python", line_numbers=True))727    console.print()728    console.print(729        "When you call [b]print()[/b], Rich [i]renders[/i] the object in to the following:\n"730    )731    fragments = list(console.render(text))732    console.print(fragments)733    console.print()734    console.print("The Segments are then processed to produce the following output:\n")735    console.print(text)736    console.print(737        "\nYou will only need to know this if you are implementing your own Rich renderables."738    )739 
codekingpro/portable-devtools · Team Ai